Spaces:
Sleeping
Sleeping
File size: 1,391 Bytes
f9d767c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 |
from langchain_community.vectorstores import FAISS
from langchain.docstore.document import Document
from langchain.embeddings.base import Embeddings
from google import genai
import os
# Make sure your environment variable GOOGLE_API_KEY is set
API_KEY = os.getenv("GOOGLE_API_KEY")
if not API_KEY:
raise ValueError("Missing GOOGLE_API_KEY environment variable!")
# Initialize client with API key
client = genai.Client(api_key=API_KEY)
class GeminiEmbeddings(Embeddings):
"""LangChain wrapper for Google Gemini embeddings"""
def embed_documents(self, texts):
if not texts:
return []
response = client.models.embed_content(
model="gemini-embedding-001",
contents=texts
)
# Each response.embeddings[i].values is a list of floats
return [e.values for e in response.embeddings]
def embed_query(self, text):
response = client.models.embed_content(
model="gemini-embedding-001",
contents=[text]
)
return response.embeddings[0].values
def create_vector_store(texts):
docs = [Document(page_content=t) for t in texts if t.strip()]
if not docs:
return None
embeddings = GeminiEmbeddings()
vectorstore = FAISS.from_texts(
texts=[d.page_content for d in docs],
embedding=embeddings
)
return vectorstore
|