admin.py

import chromadb
from chromadb.config import Settings
from fastapi import APIRouter
import pinecone
import os


router = APIRouter()


# chroma_settings = Settings(
#     anonymized_telemetry=False,
#     persist_directory="db",
#     allow_reset=False,
#     is_persistent=True,
# )
# client = chromadb.Client(chroma_settings)

# pinecone.init(
#     api_key=os.environ.get("PINECONE_API_KEY"),
#     environment=os.environ.get("PINECONE_ENV"))

@router.get("/api/v1/admin/collections")
async def get_all_collections():
    # Currently only works for ChromaDB but can be extended easily
    # for other vector stores as well
    # collections = pinecone.list_collections()
    # responses = [c.dict() for c in collections]
    # return responses
    return None


# TODO(deshraj): Add pagination and make this endpoint agnostic to the vector store
# @router.get("/api/v1/admin/collections/chromadb/{collection_name}")
# async def get_collection_details(collection_name: str):
#     collection = pinecone.describe_collection(collection_name)
#     collection_data = collection.get()
#     metadatas, documents = collection_data['metadatas'], collection_data['documents']
#     collated_data = []
#     for i in zip(metadatas, documents):
#         collated_data.append({
#             "metadata": i[0],
#             "document": i[1]
#         })
#     response = {"details": collection.dict(), "data": collated_data}
#     return response

api.py

from embedchain import Pipeline
from fastapi import APIRouter, Query, responses
from pydantic import BaseModel

router = APIRouter()

# App config using OpenAI gpt-3.5-turbo-1106 as LLM
app_config = {
    "app": {
        "config": {
            "id": "my-ai-api",
            "collect_metrics": False,
        }
    },
    "llm": {
        "provider": "openai",
        "config": {
            "model": "gpt-3.5-turbo-1106",
        }
    },
    # "vectordb": {
    #     "provider": "pinecone",
    #     "config": {
    #         "metric": "cosine",
    #         "vector_dimension": 1536,
    #         "collection_name": "embedchain-chat"
    #     }
    # }
}

# Uncomment this configuration to use Mistral as LLM
# app_config = {
#     "app": {
#         "config": {
#             "name": "embedchain-opensource-app"
#         }
#     },
#     "llm": {
#         "provider": "huggingface",
#         "config": {
#             "model": "mistralai/Mixtral-8x7B-Instruct-v0.1",
#             "temperature": 0.1,
#             "max_tokens": 250,
#             "top_p": 0.1
#         }
#     },
#     "embedder": {
#         "provider": "huggingface",
#         "config": {
#             "model": "sentence-transformers/all-mpnet-base-v2"
#         }
#     }
# }


ec_app = Pipeline.from_config(config=app_config)


class SourceModel(BaseModel):
    source: str


class QuestionModel(BaseModel):
    question: str
    session_id: str


@router.post("/api/v1/add")
async def add_source(source_model: SourceModel):
    """
    Adds a new source to the Embedchain app.
    Expects a JSON with a "source" key.
    """
    source = source_model.source
    try:
        # Add the source to a specific namespace
        
        # If I want to use the video title as the namespace name, I'd have to parse the YT link here
        # and extract the title of the video
        ec_app.add(source)
        return {"message": f"Source '{source}' added successfully."}
    except Exception as e:
        response = f"An error occurred: Error message: {str(e)}. Contact Embedchain founders on Slack: https://embedchain.com/slack or Discord: https://embedchain.com/discord"  # noqa:E501
        return {"message": response}


@router.get("/api/v1/chat")
async def handle_chat(query: str, namespace: str, session_id: str = Query(None)):
    """
    Handles a chat request to the Embedchain app.
    Accepts 'query' and 'session_id' as query parameters.
    """
    print(f"Query: {query}, Session ID: {session_id}, Namespace: {namespace}")
    try:
        response = ec_app.chat(query, session_id=session_id)
    except Exception as e:
        response = f"An error occurred: Error message: {str(e)}. Contact Embedchain founders on Slack: https://embedchain.com/slack or Discord: https://embedchain.com/discord"  # noqa:E501
    return {"response": response}


@router.get("/")
async def root():
    return responses.RedirectResponse(url="/docs")