embedding.py

# pyrefly: ignore [missing-import]
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)