Update retrieve.py

retrieve.py

@@ -1,9 +1,22 @@
+import os
 from langchain_community.vectorstores import Qdrant
 from langchain_core.output_parsers import StrOutputParser
 from langchain_core.runnables import RunnablePassthrough
 from langchain.prompts import PromptTemplate
 from utils import setup_openai_embeddings,setup_qdrant_client,openai_llm,format_document_metadata
+from dotenv import load_dotenv
+from langfuse.callback import CallbackHandler
+load_dotenv()
 
+os.environ["LANGFUSE_PUBLIC_KEY"] = os.getenv("LANGFUSE_PUBLIC_KEY")
+os.environ["LANGFUSE_SECRET_KEY"] = os.getenv("LANGFUSE_SECRET_KEY")
+os.environ["LANGFUSE_HOST"] = os.getenv("LANGFUSE_HOST") 
+
+
+ 
+langfuse_handler = CallbackHandler()
+
+ 
 
 
 def retrieve_documents(query, api_key, qdrant_url, qdrant_api_key):
@@ -41,7 +54,8 @@ def retrieve_documents(query, api_key, qdrant_url, qdrant_api_key):
         | llm
         | StrOutputParser()
     )
-    return rag_chain.invoke(query)
+    return rag_chain.invoke(query,{"callbacks":[langfuse_handler]})
+
 
 def retrieve_documents_from_collection(query, api_key, qdrant_url, qdrant_api_key, collection_name):
     """Retrieve documents based on the specified query from a specific collection."""
@@ -78,7 +92,7 @@ def retrieve_documents_from_collection(query, api_key, qdrant_url, qdrant_api_ke
         | llm
         | StrOutputParser()
     )
-    return rag_chain.invoke(query)
+    return rag_chain.invoke(query,{"callbacks":[langfuse_handler]})
 
 def delete_collection(collection_name, qdrant_url, qdrant_api_key):
     """Delete a Qdrant collection."""