Update app.py

app.py

@@ -10,10 +10,20 @@ from anthropic import Anthropic
 import pymongo
 from dotenv import load_dotenv
 import fitz  # PyMuPDF
+import voyageai
+from pinecone.grpc import PineconeGRPC as Pinecone
+from pinecone import ServerlessSpec
+from pinecone import Index
 
 # Load environment variables
 load_dotenv()
 
+# Initialize VoyageAI constants
+VOYAGEAI_BATCH_SIZE = 128
+
+# Initialize Pinecone
+PINECONE_ID = "intratalent-v2"
+
 # Initialize MongoDB client
 MONGO_URI = os.getenv('MONGO_URI')
 mongo_client = pymongo.MongoClient(MONGO_URI)
@@ -94,6 +104,62 @@ def extract_info_with_claude(resume_text: str) -> str:
     
     return extracted_info
 
+def get_pinecone_index(database_id: str) -> Index:
+    # initialize connection to pinecone
+    pc = Pinecone(api_key=os.getenv('PINECONE_API_KEY'))
+        
+    # if the index does not exist, we create it
+    if not pc.has_index(database_id):
+        pc.create_index(
+            database_id,
+            dimension=shape[1],
+            spec=ServerlessSpec(
+                cloud='aws',
+                region='us-east-1'
+            ),
+            metric='cosine'
+        )
+    
+    # connect to index
+    index = pc.Index(index_name)
+
+def add_to_voyage(person_name: str, person_projects: list) -> None:
+    embeds = []
+    metas = []
+    ids = []
+    index = get_pinecone_index(PINECONE_ID)
+    vo = voyageai.Client(api_key=os.getenv('VOYAGEAI_API_KEY'))
+    
+    for i in range(len(person_projects)):
+        # Get the ith project
+        project = person_projects[i]
+
+        # Embed the description
+        embed = vo.embed(
+            texts=project["description"],
+            model='voyage-3-lite',
+            truncation=False
+        ).embeddings[0]
+        embeds.append(embed)
+
+        # Create metadata using person's name + project name
+        meta = f"{person_name} {project["name"]}"
+        metas.append(meta)
+
+        # Give it a unique id
+        id = i
+        ids.append(i)
+
+    # create list of (id, vector, metadata) tuples to be upserted
+    to_upsert = list(zip(ids, embeds, meta))
+    
+    for i in range(0, shape[0], VOYAGEAI_BATCH_SIZE):
+        i_end = min(i+VOYAGEAI_BATCH_SIZE, shape[0])
+        index.upsert(vectors=to_upsert[i:i_end])
+    
+    # let's view the index statistics
+    st.write(index.describe_index_stats())
+
 def parse_resume(uploaded_file: UploadedFile) -> Tuple[str, List[Dict]]:
     """Parse a resume file and return name and projects."""
     try:
@@ -135,8 +201,8 @@ def parse_resume(uploaded_file: UploadedFile) -> Tuple[str, List[Dict]]:
             "projects": projects,
             "full_content": resume_text
         }
-        resume_collection.insert_one(resume_data)
-        st.write("💾 Stored data in MongoDB")
+        add_to_voyage(name, projects)
+        st.write("💾 Stored data in VoyageAI")
         
         return name, projects