update

main.py

@@ -1,19 +1,26 @@
+import time
 from fastapi import FastAPI
-from service import VectorStore
-
-PROJECT_NAME = "filecoin"
+from service.vector_store import vector_store
+from pydantic import BaseModel
+from fastapi.responses import StreamingResponse
 
 app = FastAPI()
-vector_store = VectorStore(PROJECT_NAME)
 
 @app.get("/")
 def read_root():
     return {"Hello": "World!"}
 
-@app.get("/vector/search")
-def vector_search(inputs):
+@app.get("/api/search")
+def vector_search(inputs: str):
     return vector_store.search(inputs)
 
-@app.get("/vector/marginal_search")
-def vector_marginal_search(inputs):
-    return vector_store.marginal_search(inputs)
+
+
+async def fake_video_streamer():
+    for i in range(10):
+        yield "some fake video bytes"
+        time.sleep(0.5)
+
+@app.get("/api/answer")
+async def answer(inputs: str):
+    return StreamingResponse(fake_video_streamer())

run.sh

@@ -0,0 +1 @@
+uvicorn main:app

service/vector_store.py

@@ -9,45 +9,42 @@ from langchain.vectorstores.utils import DistanceStrategy, maximal_marginal_rele
 import numpy as np
 import json
 import logging
+import uuid
+from langchain.utils.iter import batch_iterate
+try:
+    from script import export
+except:
+    pass
 
 logger = logging.getLogger(__name__)
 
-
 class VectorStore(Pinecone):
+    REQUEST_TIMEOUT=10
+    INDEX_NAME = "jarvis"
+    NAMESPACE = "filecoin"
 
-    def __init__(self,index_name) -> None:
+    def __init__(self) -> None:
         pinecone.init(
-            api_key=os.getenv("PINECONE_API_KEY"),  # find at app.pinecone.io
-            environment=os.getenv("PINECONE_ENV"),  # next to api key in console
+            api_key=os.getenv("PINECONE_API_KEY"),  
+            environment=os.getenv("PINECONE_ENV"),  
         )
-        self.index_name = index_name
+
         self.dims = 1536
-        index = pinecone.Index(self.index_name)
-        super().__init__(index, OpenAIEmbeddings(), "text")  
+        index = pinecone.Index(self.INDEX_NAME)
+        super().__init__(index, OpenAIEmbeddings(),"text")  
 
     def add_docs(self,docs):
-        if self.index_name not in pinecone.list_indexes():
-            pinecone.create_index(name=self.index_name, metric="cosine", dimension=self.dims)
-        for doc in docs:
-            metadata = doc.metadata
-            for k in metadata:
-                if metadata.get(k) is None:
-                    metadata[k] = 'unknown'
-        Pinecone.from_documents(docs, self.embeddings, index_name=self.index_name)
+        if self.INDEX_NAME not in pinecone.list_indexes():
+            pinecone.create_index(name=self.INDEX_NAME, metric="cosine", dimension=self.dims)
+        Pinecone.from_documents(docs, self.embeddings, index_name=self.INDEX_NAME)
 
-    def search(self,query,ret_json=True):
-        docs = self.similarity_search(query)
-        if ret_json:
-            return self.jsonfy(docs)
-        else:
-            return docs
+    # @timing
+    def search(self,query):
+        return self.similarity_search(query)
 
-    def marginal_search(self,query,ret_json=True):
-        docs = self.max_marginal_relevance_search(query)
-        if ret_json:
-            return self.jsonfy(docs)
-        else:
-            return docs
+    # @timing
+    def marginal_search(self,query,k=5):
+        return self.max_marginal_relevance_search(query,k=k)
         
     def jsonfy(self,docs):
         docs = [doc.dict() for doc in docs]
@@ -73,6 +70,7 @@ class VectorStore(Pinecone):
             include_metadata=True,
             namespace=namespace,
             filter=filter,
+            _request_timeout=self.REQUEST_TIMEOUT
         )
         for res in results["matches"]:
             metadata = res["metadata"]
@@ -80,7 +78,7 @@ class VectorStore(Pinecone):
                 text = metadata.pop(self._text_key)
                 score = res["score"]
                 metadata['score'] = score
-                print(f"metadata {metadata}")
+                # print(f"metadata {metadata}")
                 docs.append((Document(page_content=text, metadata=metadata), score))
             else:
                 logger.warning(
@@ -91,7 +89,7 @@ class VectorStore(Pinecone):
     def max_marginal_relevance_search_by_vector(
         self,
         embedding: List[float],
-        k: int = 4,
+        k: int = 5,
         fetch_k: int = 20,
         lambda_mult: float = 0.5,
         filter: Optional[dict] = None,
@@ -123,6 +121,7 @@ class VectorStore(Pinecone):
             include_metadata=True,
             namespace=namespace,
             filter=filter,
+            _request_timeout=self.REQUEST_TIMEOUT
         )
         mmr_selected = maximal_marginal_relevance(
             np.array([embedding], dtype=np.float32),
@@ -140,4 +139,6 @@ class VectorStore(Pinecone):
         return [
             Document(page_content=metadata.pop((self._text_key)), metadata=metadata)
             for metadata in selected
-        ]
+        ]
+
+vector_store = VectorStore()