Shifted to Pinecone

src/graph.py

@@ -4,6 +4,7 @@ from src.websearch import *
 from src.llm import *
 from langchain.schema import Document, AIMessage, HumanMessage, SystemMessage
 from typing import Annotated
+from langchain_community.vectorstores import Pinecone as LangchainPinecone
 
 from typing_extensions import TypedDict
 
@@ -75,7 +76,7 @@ def retrieve(state):
     chat_context = [serialize_messages(chat_context) for chat_context in chat_context]
     chat_context = "\n".join([d.content if hasattr(d, 'content') else d["content"] for d in chat_context])
 
-    documents = retriever.invoke("Question: " + question + " Chat Context: " + chat_context)
+    documents = retriever.get_relevant_documents("Question: " + question + " Chat Context: " + chat_context)
     print("---RETRIEVED---")
     return {"documents": documents}
 

src/index.py

@@ -4,11 +4,12 @@ import os
 from dotenv import load_dotenv
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_community.document_loaders import PyPDFLoader
-from langchain_community.vectorstores import Chroma
 from langchain_openai import AzureOpenAIEmbeddings, AzureChatOpenAI
 import time
 from tenacity import retry, stop_after_attempt, wait_exponential
-from tqdm import tqdm  # Add this import for progress bar
+from tqdm import tqdm
+from pinecone import Pinecone, ServerlessSpec
+from langchain_community.vectorstores import Pinecone as LangchainPinecone
 
 # Load environment variables
 load_dotenv()
@@ -25,6 +26,11 @@ try:
     api_key = os.getenv("API_KEY")
     api_version = os.getenv("API_VERSION")
     
+    # Pinecone environment variables
+    PINECONE_API_KEY = os.getenv("PINECONE_API_KEY")
+    PINECONE_ENVIRONMENT = os.getenv("PINECONE_ENVIRONMENT")
+    PINECONE_INDEX_NAME = os.getenv("PINECONE_INDEX_NAME")
+    
     print("Environment variables loaded successfully.")
 except Exception as e:
     print(f"Error loading environment variables: {e}")
@@ -49,9 +55,12 @@ except Exception as e:
     print(f"Error setting up Azure OpenAI: {e}")
     sys.exit(1)
 
-# Function to check if vector store exists
-def vector_store_exists(persist_directory):
-    return os.path.exists(persist_directory) and len(os.listdir(persist_directory)) > 0
+# Initialize Pinecone
+pc = Pinecone(api_key=PINECONE_API_KEY, environment=PINECONE_ENVIRONMENT)
+
+# Function to check if Pinecone index exists
+def pinecone_index_exists(index_name):
+    return index_name in pc.list_indexes().names()
 
 # Load and process documents
 try:
@@ -61,7 +70,7 @@ try:
 
     print("Splitting documents...")
     text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(
-        chunk_size=300, chunk_overlap=100
+        chunk_size=400, chunk_overlap=100
     )
     doc_splits = text_splitter.split_documents(docs)
     print(f"Documents split into {len(doc_splits)} chunks.")
@@ -70,42 +79,58 @@ except Exception as e:
     sys.exit(1)
 
 @retry(stop=stop_after_attempt(5), wait=wait_exponential(multiplier=1, min=4, max=10))
-def create_vector_store_batch(persist_directory, documents, embedding, batch_size=50):
-    vectorstore = None
+def create_pinecone_index(index_name, dimension, spec):
+    try:
+        if not pinecone_index_exists(index_name):
+            print(f"Creating new Pinecone index: {index_name}")
+            pc.create_index(
+                name=index_name,
+                dimension=dimension,
+                metric='cosine',
+                spec=spec
+            )
+        print(f"Connecting to Pinecone index: {index_name}")
+        return pc.Index(index_name)
+    except Exception as e:
+        print(f"Error creating/connecting to Pinecone index: {e}")
+        raise
+
+@retry(stop=stop_after_attempt(5), wait=wait_exponential(multiplier=1, min=4, max=10))
+def upsert_to_pinecone(index, documents, embedding, batch_size=50):
     for i in tqdm(range(0, len(documents), batch_size), desc="Processing batches"):
         batch = documents[i:i+batch_size]
-        if vectorstore is None:
-            vectorstore = Chroma.from_documents(
-                documents=batch,
-                collection_name="rag-chroma",
-                embedding=embedding,
-                persist_directory=persist_directory
-            )
-        else:
-            vectorstore.add_documents(batch)
+        ids = [str(j) for j in range(i, min(i+batch_size, len(documents)))]
+        embeds = embedding.embed_documents([doc.page_content for doc in batch])
+        metadata = [{"text": doc.page_content} for doc in batch]
+        to_upsert = list(zip(ids, embeds, metadata))
+        index.upsert(vectors=to_upsert)
         time.sleep(1)  # Add a small delay between batches
-    return vectorstore
 
-# Create or load vector store
+# Create or load Pinecone index
 try:
-    persist_directory = './vectordb'
-    if not vector_store_exists(persist_directory):
-        print("Creating new vector store...")
-        vectorstore = create_vector_store_batch(persist_directory, doc_splits, embd)
-        print("New vector store created and populated.")
+    print("Setting up Pinecone index...")
+    dimension = 1536  # Dimension for Azure OpenAI embeddings
+    pinecone_index = create_pinecone_index(PINECONE_INDEX_NAME, dimension, spec=ServerlessSpec(cloud='aws', region='us-east-1'))
+    
+    print("Checking index statistics...")
+    index_stats = pinecone_index.describe_index_stats()
+    print(f"Index stats: {index_stats}")
+    
+    if index_stats['total_vector_count'] == 0:
+        print("Upserting documents to Pinecone...")
+        upsert_to_pinecone(pinecone_index, doc_splits, embd)
+        print("Documents upserted to Pinecone successfully.")
     else:
-        print("Loading existing vector store...")
-        vectorstore = Chroma(
-            persist_directory=persist_directory,
-            embedding_function=embd,
-            collection_name="rag-chroma"
-        )
-        print("Existing vector store loaded.")
+        print("Pinecone index already populated.")
 
-    retriever = vectorstore.as_retriever(search_kwargs={"k": 10, "score_threshold": 0.6}, search_type="similarity_score_threshold")
+    print("Creating LangChain vectorstore...")
+    vectorstore = LangchainPinecone(pinecone_index, embd.embed_query, "text")
+    retriever = vectorstore.as_retriever(search_kwargs={"k": 10})
     print("Retriever set up successfully.")
 except Exception as e:
-    print(f"Error with vector store operations: {e}")
+    print(f"Error with Pinecone operations: {e}")
+    import traceback
+    traceback.print_exc()
     sys.exit(1)
 
 print("Index setup completed successfully.")