File size: 749 Bytes
8953dfc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 |
import os
import google.generativeai as genai
from chromadb.api.types import Documents, Embeddings
from chromadb import EmbeddingFunction
from dotenv import load_dotenv
load_dotenv()
gemini_api_key = os.environ["GEMINI_API_KEY"]
class GeminiEmbeddingFunction(EmbeddingFunction):
"""
Custom embedding function using Gemini AI API.
"""
def __call__(self, input: Documents) -> Embeddings:
if not gemini_api_key:
raise ValueError("Gemini API Key not provided. Please set GEMINI_API_KEY as an environment variable.")
genai.configure(api_key=gemini_api_key)
model = "models/text-embedding-004"
return genai.embed_content(model=model, content=input, task_type="retrieval_document")["embedding"] |