|
|
| from langchain_google_genai import GoogleGenerativeAIEmbeddings
|
| from config import get_settings
|
| from rag.repository_loader import CodeDocument
|
|
|
|
|
| def get_embedding_client() -> GoogleGenerativeAIEmbeddings:
|
| settings = get_settings()
|
|
|
| return GoogleGenerativeAIEmbeddings(
|
| model=settings.embedding_model,
|
| google_api_key=settings.gemini_api_key,
|
| )
|
|
|
|
|
| def embed_document(chunks: list[dict]) -> list[dict]:
|
| embeddings = get_embedding_client()
|
| texts = [chunk["content"] for chunk in chunks]
|
| vectors = embeddings.embed_documents(texts)
|
|
|
| for chunk, vector in zip(chunks, vectors):
|
| chunk["embedding"] = vector
|
| return chunks
|
|
|
|
|
| def embed_query(query: str) -> list[float]:
|
| embeddings = get_embedding_client()
|
| return embeddings.embed_query(query)
|
|
|
|
|
|
|
