graph_rag_server/model/tracked_openai_embeddings.py

from langchain_openai import OpenAIEmbeddings
from utils.decorators.langfuse_trace_embedding import langfuse_trace_embedding

class TrackedOpenAIEmbeddings(OpenAIEmbeddings):
    
    @langfuse_trace_embedding(name="openai-embed-query")
    def embed_query(self, text: str):
        return super().embed_query(text)

    @langfuse_trace_embedding(name="openai-embed-documents")
    def embed_documents(self, texts: list[str]):
        return super().embed_documents(texts)