refactor

api/db/vector_store.py

@@ -1,4 +1,5 @@
 from abc import abstractmethod
+from functools import cache
 import os
 from qdrant_client import QdrantClient
 from langchain.embeddings.sentence_transformer import SentenceTransformerEmbeddings
@@ -7,7 +8,7 @@ from qdrant_client.models import VectorParams, Distance
 from db.embedding import Embedding, EMBEDDINGS
 
 
-class ToyVectorStore:
+class Store:
 
     @staticmethod
     def get_embedding():
@@ -17,13 +18,14 @@ class ToyVectorStore:
         return EMBEDDINGS[embedding]
     
     @staticmethod
+    @cache
     def get_instance():
         vector_store = os.getenv("STORE")
 
         if vector_store == "ELASTIC":
-            return ElasticVectorStore(ToyVectorStore.get_embedding())
+            return ElasticVectorStore(Store.get_embedding())
         elif vector_store == "QDRANT":
-            return QdrantVectorStore(ToyVectorStore.get_embedding())
+            return QdrantVectorStore(Store.get_embedding())
         else:
             raise ValueError(f"Invalid vector store {vector_store}")
     
@@ -47,7 +49,14 @@ class ToyVectorStore:
         """
         pass
 
-class ElasticVectorStore(ToyVectorStore):
+    @abstractmethod
+    def list_collections(self) -> list[dict]:
+        """
+        Return a list of collections in the vecot store.
+        """
+        pass
+
+class ElasticVectorStore(Store):
     def __init__(self, embeddings):
         super().__init__(embeddings)
 
@@ -59,8 +68,11 @@ class ElasticVectorStore(ToyVectorStore):
         store = self.get_collection(collection)
         store.create_index(store.client,collection, dict())
 
+    def list_collections(self) -> list[dict]:
+        #TODO: not impelented
+        return []
 
-class QdrantVectorStore(ToyVectorStore):
+class QdrantVectorStore(Store):
 
     def __init__(self, embeddings):
         super().__init__(embeddings)
@@ -75,4 +87,9 @@ class QdrantVectorStore(ToyVectorStore):
         self.client.create_collection(collection_name=collection, 
                         vectors_config=VectorParams(size=self.embedding.dimension, 
                                                     distance=Distance.COSINE))
+    
+    def list_collections(self) -> list[dict]:
+        """ return a list of collections.
+        """
+        return [ c for i,c in enumerate(self.client.get_collections().collections)]
     

api/document_parsing.py

@@ -0,0 +1,35 @@
+from typing import Annotated
+
+from fastapi import APIRouter, UploadFile, File, Body
+from langchain.schema import Document
+import io
+import os
+from pypdf import PdfReader
+from langchain.text_splitter import SentenceTransformersTokenTextSplitter
+from db.vector_store import Store
+
+async def generate_documents(file: UploadFile, file_name: str):
+    num=0
+    async for txts in convert_documents(file):
+        num += 1
+        for txt in txts:
+            document = Document(page_content=txt,metadata={"file": file_name, "page": num})
+            yield document
+ 
+async def convert_documents(file: UploadFile):
+    splitter = SentenceTransformersTokenTextSplitter(chunk_overlap=0)
+
+    #parse pdf document
+    if file.content_type == 'application/pdf':
+        content = await file.read()
+        pdf_reader = PdfReader(io.BytesIO(content))
+        try:
+            for page in pdf_reader.pages:
+                yield splitter.split_text(page.extract_text())
+        except Exception as e:
+            print(f"Exception {e}")
+    elif "text" in file.content_type:
+        content = await file.read()
+        yield splitter.split_text(content.decode("utf-8"))
+    else:
+        return

api/main.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 from fastapi import FastAPI
-from routes import embeddings, search, admin
+from routes import search, upload
 from fastapi.middleware import Middleware
 from fastapi.middleware.cors import CORSMiddleware
 from datetime import datetime
@@ -16,9 +16,8 @@ logger.addHandler(handler)
 
 # Create the FastAPI instance
 app = FastAPI()
-app.include_router(embeddings.router)
 app.include_router(search.router)
-app.include_router(admin.router)
+app.include_router(upload.router)
 app.exception_handler(generic_exception_handler)
 
 app.add_middleware(CORSMiddleware, allow_origins = ["*"],

api/routes/admin.py

@@ -1,16 +0,0 @@
-#This is to init the vector store
-
-from typing import Annotated
-
-from fastapi import APIRouter, Body
-from db.vector_store import ToyVectorStore
-
-router = APIRouter()
-
-@router.put("/admin/v1/db")
-async def recreate_collection(name: Annotated[str, Body(embed=True)]):
-    """ `name` of the collection to be created.
-    If one exits, delete and recreate.
-    """
-    print(f"creating collection {name} in db")
-    return ToyVectorStore.get_instance().create_collection(name)

api/routes/embeddings.py

@@ -1,15 +0,0 @@
-from fastapi import APIRouter, UploadFile, File
-import openai
-import io
-import os
-from pypdf import PdfReader
-
-router = APIRouter()
-
-openai.api_key = os.getenv("OPENAI_API_KEY")
-
-@router.post("/v1/embeddings")
-async def embed_doc(file: UploadFile = File(...)):    
-    #for now just truncate based on length of words
-    content = await file.read()
-    return openai.Embedding.create(input = content.decode("utf-8"), model = "text-embedding-ada-002")

api/routes/search.py

@@ -10,65 +10,28 @@ from langchain.schema import Document
 from langchain.chains.question_answering import load_qa_chain
 from langchain.llms import OpenAI
 from langchain.text_splitter import SentenceTransformersTokenTextSplitter
-from db.vector_store import ToyVectorStore
+from db.vector_store import Store
 
 router = APIRouter()
 _chain = load_qa_chain(OpenAI(temperature=0), chain_type="stuff", verbose=True)
 
-@router.post("/v1/docs")
-async def create_or_update(name: Annotated[str, Body()], file_name: Annotated[str, Body()], file: UploadFile = File(...)):
-    """Create or update an existing collection with information from the file 
-    `name` of the collection
-    `file` to upload.
-    `fileName` name of the file.
-    """
-    
-    _db = ToyVectorStore.get_instance().get_collection(name)
-    if not _db:
-        #todo. fix this to create a collection, may be.
-        return JSONResponse(status_code=404, content={})
-
-    async for doc in generate_documents(file, file_name):
-        print(doc)
-        _db.add_documents([doc])
-    #todo return something sensible
-    return JSONResponse(status_code=200, content={"name": name})
 
-@router.get("/v1/doc/{name}/answer")
+@router.get("/v1/docs/{name}/answer")
 async def answer(name: str, query: str):
-    """ Answer a question from the collection
-    `name` of the collection.
+    """ Answer a question from the doc
+    `name` of the doc.
     `query` to be answered.
     """
-    _db = ToyVectorStore.get_instance().get_collection(name)
+    _db = Store.get_instance().get_collection(name)
     print(query)
     docs = _db.similarity_search_with_score(query=query)
     print(docs)
     answer = _chain.run(input_documents=[tup[0] for tup in docs], question=query)
     return JSONResponse(status_code=200, content={"answer": answer, "file_score": [[f"{d[0].metadata['file']} : {d[0].metadata['page']}", d[1]] for d in docs]}) 
 
-async def generate_documents(file: UploadFile, file_name: str):
-    num=0
-    async for txts in convert_documents(file):
-        num += 1
-        for txt in txts:
-            document = Document(page_content=txt,metadata={"file": file_name, "page": num})
-            yield document
- 
-async def convert_documents(file: UploadFile):
-    splitter = SentenceTransformersTokenTextSplitter(chunk_overlap=0)
 
-    #parse pdf document
-    if file.content_type == 'application/pdf':
-        content = await file.read()
-        pdf_reader = PdfReader(io.BytesIO(content))
-        try:
-            for page in pdf_reader.pages:
-                yield splitter.split_text(page.extract_text())
-        except Exception as e:
-            print(f"Exception {e}")
-    elif "text" in file.content_type:
-        content = await file.read()
-        yield splitter.split_text(content.decode("utf-8"))
-    else:
-        return
+@router.get("/v1/docs")
+async def list() -> list[dict]:
+    """ List all the docs. 
+    """
+    return Store.get_instance().list_collections()

api/routes/upload.py

@@ -0,0 +1,39 @@
+#This is to init the vector store
+
+from typing import Annotated
+
+from db.vector_store import Store
+from document_parsing import generate_documents
+
+from fastapi import APIRouter, Body
+from fastapi.responses import JSONResponse
+from fastapi import APIRouter, UploadFile, File, Body
+
+router = APIRouter()
+
+@router.put("/v1/docs")
+async def recreate_collection(name: Annotated[str, Body(embed=True)]):
+    """ `name` of the doc to be created.
+    If one exits, delete and recreate.
+    """
+    print(f"creating collection {name} in db")
+    return Store.get_instance().create_collection(name)
+
+@router.post("/v1/docs")
+async def update(name: Annotated[str, Body()], file_name: Annotated[str, Body()], file: UploadFile = File(...)):
+    """Update an existing document with information from the file.
+    If one doesn't exist with name, it creates a new document to update. 
+    `name` of the collection
+    `file` to upload.
+    `fileName` name of the file. This is used for metadata purposes only.
+    """
+    
+    _db = Store.get_instance().get_collection(name)
+    if not _db:
+        return JSONResponse(status_code=404, content={})
+
+    async for doc in generate_documents(file, file_name):
+        print(doc)
+        _db.add_documents([doc])
+    return JSONResponse(status_code=200, content={"name": name})
+