Introduce elastic search as backened store

Idea is to see if it works any better if we use
elastic search based lookup.

In specific if we use elastic's sparse encoder

api/db/vector_store.py

@@ -1,12 +1,60 @@
+from abc import abstractmethod
 import os
 from qdrant_client import QdrantClient
 from langchain.embeddings import OpenAIEmbeddings
-from langchain.vectorstores import Qdrant
+from langchain.vectorstores import Qdrant, ElasticVectorSearch, VectorStore
+from qdrant_client.models import VectorParams, Distance
 
 embeddings = OpenAIEmbeddings()
-client = QdrantClient(url=os.getenv("QDRANT_URL"),
-                                    api_key=os.getenv("QDRANT_API_KEY"))
 
-def get_instance(collection: str = "test"):
-    return Qdrant(client=client,collection_name=collection,embeddings=embeddings)
+class ToyVectorStore:
 
+    @staticmethod
+    def get_instance():
+        vector_store = os.getenv("STORE")
+        if vector_store == "ELASTIC":
+            return ElasticVectorStore()
+        elif vector_store == "QDRANT":
+            return QdrantVectorStore()
+        else:
+            raise ValueError(f"Invalid vector store {vector_store}")
+            
+    @abstractmethod
+    def get_collection(self, collection: str="test") -> VectorStore:
+        """
+        get an instance of vector store
+        of collection
+        """
+        pass
+    
+    @abstractmethod
+    def create_collection(self, collection: str) -> None:
+        """
+        create an instance of vector store
+        with collection name
+        """
+        pass
+
+class ElasticVectorStore(ToyVectorStore):
+    def get_collection(self, collection:str) -> ElasticVectorSearch:
+        return ElasticVectorSearch(elasticsearch_url= os.getenv("ES_URL"),
+                               index_name= collection, embedding=embeddings)
+
+    def create_collection(self, collection: str) -> None:
+        store = self.get_collection(collection)
+        store.create_index(store.client,collection, dict())
+
+
+class QdrantVectorStore(ToyVectorStore):
+
+    def __init__(self):
+        self.client = QdrantClient(url=os.getenv("QDRANT_URL"),
+                                        api_key=os.getenv("QDRANT_API_KEY"))
+
+    def get_collection(self, collection: str) -> Qdrant:  
+        return Qdrant(client=self.client,collection_name=collection,embeddings=embeddings)
+
+    def create_collection(self, collection: str) -> None:
+        self.client.create_collection(collection_name=collection, 
+                        vectors_config=VectorParams(size=1536, distance=Distance.COSINE))
+    

api/routes/admin.py

@@ -2,9 +2,8 @@
 
 from typing import Annotated
 
-from qdrant_client.models import VectorParams, Distance
 from fastapi import APIRouter, Body
-from db import vector_store
+from db.vector_store import ToyVectorStore
 
 router = APIRouter()
 
@@ -14,5 +13,4 @@ async def recreate_collection(name: Annotated[str, Body(embed=True)]):
     If one exits, delete and recreate.
     """
     print(f"creating collection {name} in db")
-    return vector_store.client.recreate_collection(collection_name=name, 
-                                            vectors_config=VectorParams(size=1536, distance=Distance.COSINE))
+    return ToyVectorStore.get_instance().create_collection(name)

api/routes/search.py

@@ -11,22 +11,24 @@ from langchain.vectorstores import Qdrant
 from langchain.schema import Document
 from langchain.chains.question_answering import load_qa_chain
 from langchain.llms import OpenAI
-from db import vector_store
+from db.vector_store import ToyVectorStore
 
 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 a collection 
+    """Create or update an existing collection with information from the file 
     `name` of the collection
     `file` to upload.
     `fileName` name of the file.
     """
 
-    _db = vector_store.get_instance(name)
+    _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])
@@ -39,9 +41,9 @@ async def answer(name: str, query: str):
     `name` of the collection.
     `query` to be answered.
     """
-    _db = vector_store.get_instance(name)
+    _db = ToyVectorStore.get_instance().get_collection(name)
     print(query)
-    docs = _db.similarity_search_with_relevance_scores(query=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]}) 

requirements.txt

@@ -8,4 +8,4 @@ langchain
 tiktoken
 faiss-cpu
 qdrant-client
-
+elasticsearch