Update app/spark.py

app/spark.py

@@ -23,17 +23,25 @@ from chainlit import on_message, on_chat_start
 import openai
 from langchain.callbacks import ContextCallbackHandler
 from promptwatch import PromptWatch
+import os
+from pinecone import Pinecone, ServerlessSpec
+from langchain_openai import OpenAIEmbeddings
+from langchain_pinecone import PineconeVectorStore
+
+pc = Pinecone(
+        api_key=os.environ.get("PINECONE_API_KEY")
+    )
 
 
-index_name = "spark"
+index_name = "sparklearn"
 
 spark = load_spark_prompt()
 query_gen_prompt = load_query_gen_prompt()
 CONDENSE_QUESTION_PROMPT = PromptTemplate.from_template(query_gen_prompt)
-pinecone.init(
-            api_key=os.environ.get("PINECONE_API_KEY"),
-            environment='us-west1-gcp',
-    )
+# pinecone.init(
+#             api_key=os.environ.get("PINECONE_API_KEY"),
+#             environment='us-west1-gcp',
+#     )
 @on_chat_start
 def init(): 
     token = os.environ["CONTEXT_TOKEN"]
@@ -43,12 +51,19 @@ def init():
     llm = ChatOpenAI(temperature=0.7, verbose=True, openai_api_key = os.environ.get("OPENAI_API_KEY"), streaming=True,
                      callbacks=[context_callback])
     memory = ConversationTokenBufferMemory(llm=llm,memory_key="chat_history", return_messages=True,input_key='question',max_token_limit=1000)
-    embeddings = CohereEmbeddings(model='embed-english-light-v2.0',cohere_api_key=os.environ.get("COHERE_API_KEY"))
+    # embeddings = CohereEmbeddings(model='embed-english-light-v2.0',cohere_api_key=os.environ.get("COHERE_API_KEY"))
+    embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
+    index = pc.Index(index_name)
 
-    docsearch = Pinecone.from_existing_index(
-    index_name=index_name, embedding=embeddings
-    )
-    retriever = docsearch.as_retriever(search_kwargs={"k": 4})
+
+
+    # docsearch = Pinecone.from_existing_index(
+    # index_name=index_name, embedding=embeddings
+    # )
+    vector_store = PineconeVectorStore(index=index, embedding=embeddings)
+
+
+    retriever = vector_store.as_retriever(search_kwargs={"k": 4}, search_type="similarity_score_threshold")
     # compressor = CohereRerank()
     # reranker = ContextualCompressionRetriever(
     #     base_compressor=compressor, base_retriever=retriever