Updated Query Classifier

langgraph_agent.py

@@ -1,11 +1,9 @@
-# Import LangGraph components
 from langgraph.graph import StateGraph, END
-from typing import TypedDict, List, Dict, Any, Annotated, Union, Literal
-import operator
-from pydantic import BaseModel, Field
+from typing import TypedDict, List, Dict, Any
 from agent import LLMChain, PromptTemplate, llm, DOCUMENT_DIR, load_documents, split_documents, CHROMA_PATH, load_vectordb, create_and_store_embeddings
 import os
-# Define state schema
+
+# state schema
 class AgentState(TypedDict):
     query: str
     previous_conversation: str
@@ -14,18 +12,36 @@ class AgentState(TypedDict):
     context: List[str]
     response: str
     
-# Define tools and nodes for the LangGraph
+
+
+# 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:
-    """Determine if the query requires RAG retrieval."""
-    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"
-    ]
+    """Updated classifier to use LLM for intent classification."""
+    
+    query = state["query"]
+    
+    # Then classify intent
+    classification_prompt = f"""
+    Answer with only 'Yes' or 'No'.
+    Classify if this query is asking about government schemes, policies, or benefits.
+    The language may not be English, So first detect the language. and understand the query.:
+    Query: {query}
+    Remember Answer with only 'Yes' or 'No'."""
     
-    state["requires_rag"] = any(keyword in query_lower for keyword in rag_keywords)
+    result = llm.predict(classification_prompt)
+    state["requires_rag"] = "yes" in result.lower()
     return state
 
 def retrieve_documents(state: AgentState) -> AgentState:
@@ -183,18 +199,12 @@ def agent_with_db():
                 "requires_rag": False,
                 "context": [],
                 "response": "",
-                # "style": style
             }
-            print("Initial state:", initial_state)
             
-            # Run the workflow
             final_state = self.workflow.invoke(initial_state)
-            print("Final state:", final_state)
             
-            # Update conversation history with response
             self.conversation_history += f"Assistant: {final_state['response']}\n"
             
-            # Return in the expected format
             return {"result": final_state["response"]}
         
     return HealthAgent(agent_workflow)

main.py

@@ -61,7 +61,6 @@ async def retrieve(request:request, url:Request):
         if origin is None:
             origin = url.headers.get('referer')
         print("origin: ", origin)
-        print("response: ", response)
         return {"response": response["result"]}
     
     except Exception as e: