enhanced_query

generate.py

@@ -1,8 +0,0 @@
-from google import genai
-from dotenv import load_dotenv
-from os import getenv
-
-load_dotenv()
-
-GEMINI_API_KEY = getenv("GEMINI_API_KEY")
-from .constants import GEMINI_API_KEY

langgraph_agent.py

@@ -12,21 +12,6 @@ class AgentState(TypedDict):
     context: List[str]
     response: str
     
-
-
-# def query_classifier(state: AgentState) -> AgentState:
-#     """Determine if the query requires RAG retrieval based on keywords.
-#       Is not continued anymore, will be removed in future."""
-#     query_lower = state["query"].lower()
-#     rag_keywords = [
-#         "scheme", "schemes", "program", "programs", "policy", "policies", 
-#         "public health engineering", "phe", "public health", "government", 
-#         "benefit", "financial", "assistance", "aid", "initiative", "yojana",
-#     ]
-    
-#     state["requires_rag"] = any(keyword in query_lower for keyword in rag_keywords)
-#     return state
-
 def query_classifier(state: AgentState) -> AgentState:
     """Updated classifier to use LLM for intent classification."""
     
@@ -44,6 +29,25 @@ def query_classifier(state: AgentState) -> AgentState:
     state["requires_rag"] = "yes" in result.lower()
     return state
 
+def enhance_query(state:AgentState) -> AgentState:
+    """Enhance the query with user data and context."""
+    previous_conversation = state.get("previous_conversation", "")
+    user_data = state.get("user_data", {})
+    query = state.get("query", "")
+    
+    query_enhancement_prompt = f"""
+    Enhance the following query with user data and previous conversation context so it uses the previous conversation and user data.
+    To be used for generating a more relevant and personalized response.
+    Previous Conversation: {previous_conversation}
+    User Data: {user_data}
+    Current Query: {query}
+    Only write the enhanced query. No other text."""
+    result = llm.predict(query_enhancement_prompt)
+    print("Enhanced query: ", result)
+    state["query"] = result
+
+    return state
+
 def retrieve_documents(state: AgentState) -> AgentState:
     """Retrieve documents from vector store if needed."""
     if state["requires_rag"]:
@@ -121,17 +125,19 @@ def create_agent_workflow():
     workflow = StateGraph(AgentState)
     
     # Add nodes
+    workflow.add_node("enhance_query", enhance_query)
     workflow.add_node("classifier", query_classifier)
     workflow.add_node("retriever", retrieve_documents)
     workflow.add_node("responder", generate_response)
     
     # Create edges
+    workflow.add_edge("enhance_query", "classifier")
     workflow.add_edge("classifier", "retriever")
     workflow.add_edge("retriever", "responder")
     workflow.add_edge("responder", END)
     
     # Set the entry point
-    workflow.set_entry_point("classifier")
+    workflow.set_entry_point("enhance_query")
     
     # Compile the graph
     return workflow.compile()

requirements.txt

@@ -273,3 +273,6 @@ xxhash==3.5.0
 yarl==1.18.3
 zipp==3.21.0
 zstandard==0.23.0
+tf-keras
+# pathway[all]
+# langchain-google-genai