Inital Commit

backend/src/rag/embeddings.py

@@ -0,0 +1,33 @@
+from google import genai
+import os
+
+def get_embedding(text, model="gemini-embedding-001"):
+    """
+    Generates an embedding for the given text using the Gemini API.
+    """
+    api_key = os.environ.get("GOOGLE_API_KEY")
+    if not api_key:
+        raise ValueError("GOOGLE_API_KEY environment variable not set")
+    
+    client = genai.Client(api_key=api_key)
+    result = client.models.embed_content(
+        model=model,
+        contents=text
+    )
+    return result.embeddings[0].values
+
+def get_embeddings_batch(texts, model="gemini-embedding-001"):
+    """
+    Generates embeddings for a list of texts.
+    """
+    api_key = os.environ.get("GOOGLE_API_KEY")
+    if not api_key:
+        raise ValueError("GOOGLE_API_KEY environment variable not set")
+        
+    client = genai.Client(api_key=api_key)
+    result = client.models.embed_content(
+        model=model,
+        contents=texts
+    )
+    # The SDK returns a list of embedding objects
+    return [emb.values for emb in result.embeddings]