Create helper/create_embeddings.py

helper/create_embeddings.py

@@ -0,0 +1,93 @@
+import pinecone
+from langchain_google_genai import GoogleGenerativeAIEmbeddings
+import os
+import uuid
+from pymongo import MongoClient
+from dotenv import load_dotenv
+load_dotenv()
+FLASH_API = os.getenv("FLASH_API")
+PINECONE_API=os.getenv("PINECONE_API")
+PINECONE_INDEX=os.getenv("PINECONE_INDEX")
+google_embeddings = GoogleGenerativeAIEmbeddings(
+    model="models/embedding-001",  # Correct model name
+    google_api_key=FLASH_API  # Your API key
+)
+
+# Initialize Pinecone instance
+pc = pinecone.Pinecone(
+    api_key=PINECONE_API  # Your Pinecone API key
+)
+
+MONGO_URI = os.getenv("MONGO_URI")
+DB_NAME = os.getenv("DB_NAME")
+COLLECTION_NAME = os.getenv("COLLECTION_NAME")
+
+mongo_client = MongoClient(MONGO_URI)
+db = mongo_client[DB_NAME]
+collection = db[COLLECTION_NAME]
+# Define the Pinecone index name (make sure it exists in your Pinecone dashboard)
+
+index = pc.Index(PINECONE_INDEX)
+
+
+def create_embedding(object_url,tags,categories):
+    try:
+        document = collection.find_one({'object_url': object_url})
+        content = document.get("description")
+        file_type = document.get("type")
+        mongo_id = str(document.get('_id'))  # Convert ObjectId to string for storage in metadata
+
+        # Generate the embedding
+        embedding = google_embeddings.embed_query(content)
+
+        # Generate a unique ID for Pinecone
+        pinecone_id = str(uuid.uuid4())
+
+        # Prepare the vector with metadata
+        vectors = [{
+            'id': pinecone_id,
+            'values': embedding,
+            'metadata': {
+                'description': content,
+                "url": object_url,
+                "filetype": file_type,
+                "mongo_id": mongo_id,
+                "tags":','.join(tags),
+                "categories":','.join(categories)# Store MongoDB ID in metadata
+            }
+        }]
+
+        # Upsert the vector to Pinecone
+        index.upsert(vectors)
+        print(f"Inserted: {object_url} in Pinecone with MongoDB ID in metadata")
+
+        # Update MongoDB document with Pinecone ID and success status
+        collection.update_one(
+            {"_id": document["_id"]},
+            {"$set": {
+                "pinecone_id": pinecone_id,
+                "successfully_embedding_created": True
+            }}
+        )
+
+        return True
+    except Exception as e:
+        print(f"Error occurred: {e}")
+
+        # Update MongoDB document with failure status
+        collection.update_one(
+            {"_id": document["_id"]},
+            {"$set": {
+                "successfully_embedding_created": False
+            }}
+        )
+
+        return False
+
+
+
+
+
+
+
+