Update agent.py

agent.py

@@ -509,25 +509,39 @@ def process_question(question):
 
 
 
-def retriever(state: MessagesState):
-    """Retriever node using similarity scores for filtering"""
-    query = state["messages"][0].content
-    results = vector_store.similarity_search_with_score(query, k=4)  # top 4 matches
-    
-    # Dynamically adjust threshold based on query complexity
-    threshold = 0.75 if "who" in query else 0.8
+from langchain.schema import HumanMessage
+
+def retriever(state: MessagesState, k: int = 4):
+    """
+    Retrieves documents from the vector store using similarity scores, 
+    applies a dynamic threshold filter, and returns updated message state.
+
+    Args:
+        state (MessagesState): Current message state including the user's query.
+        k (int): Number of top results to retrieve from the vector store.
+
+    Returns:
+        dict: Updated messages state including relevant documents or fallback message.
+    """
+    query = state["messages"][0].content.strip()
+    results = vector_store.similarity_search_with_score(query, k=k)
+
+    # Determine dynamic similarity threshold
+    if any(keyword in query.lower() for keyword in ["who", "what", "where", "when", "why", "how"]):
+        threshold = 0.75
+    else:
+        threshold = 0.8
+
     filtered = [doc for doc, score in results if score < threshold]
 
-    # Provide a default message if no documents found
     if not filtered:
-        example_msg = HumanMessage(content="No relevant documents found.")
+        response_msg = HumanMessage(content="No relevant documents found.")
     else:
         content = "\n\n".join(doc.page_content for doc in filtered)
-        example_msg = HumanMessage(
-            content=f"Here are relevant reference documents:\n\n{content}"
-        )
+        response_msg = HumanMessage(content=f"Here are relevant reference documents:\n\n{content}")
+
+    return {"messages": [sys_msg] + state["messages"] + [response_msg]}
 
-    return {"messages": [sys_msg] + state["messages"] + [example_msg]}
 
 
 
@@ -537,17 +551,29 @@ def retriever(state: MessagesState):
 def get_llm(provider: str, config: dict):
     if provider == "google":
         from langchain_google_genai import ChatGoogleGenerativeAI
-        return ChatGoogleGenerativeAI(model=config["model"], temperature=config["temperature"])
+        return ChatGoogleGenerativeAI(
+            model=config.get("model"),
+            temperature=config.get("temperature", 0.7),
+            google_api_key=config.get("api_key")  # Optional: if needed
+        )
 
     elif provider == "groq":
         from langchain_groq import ChatGroq
-        return ChatGroq(model=config["model"], temperature=config["temperature"])
+        return ChatGroq(
+            model=config.get("model"),
+            temperature=config.get("temperature", 0.7),
+            groq_api_key=config.get("api_key")  # Optional: if needed
+        )
 
     elif provider == "huggingface":
         from langchain_huggingface import ChatHuggingFace
         from langchain_huggingface import HuggingFaceEndpoint
         return ChatHuggingFace(
-            llm=HuggingFaceEndpoint(url=config["url"], temperature=config["temperature"])
+            llm=HuggingFaceEndpoint(
+                endpoint_url=config.get("url"),
+                temperature=config.get("temperature", 0.7),
+                huggingfacehub_api_token=config.get("api_key")  # Optional
+            )
         )
 
     else:
@@ -607,7 +633,7 @@ def planner(question: str, tools: list) -> tuple:
     return detected_intent, matched_tools if matched_tools else [tools[0]]
 
 
-import json
+
 
 def task_classifier(question: str) -> str:
     """
@@ -889,7 +915,7 @@ def tool_dispatcher(state: AgentState) -> AgentState:
 
 
 # Decide what to do next: if tool call → call_tool, else → end
-def call_tool(state):
+def should_call_tool(state):
     last_msg = state["messages"][-1]
     if isinstance(last_msg, AIMessage) and last_msg.tool_calls:
         return "call_tool"
@@ -904,9 +930,6 @@ class AgentState(TypedDict):
     intent: str                   # derived or predicted intent
     result: Optional[str]         # tool output, if any
 
-builder.add_node("call_tool", tool_dispatcher)
-
-
 
 
 
@@ -1036,14 +1059,7 @@ model_config = {
     "huggingfacehub_api_token": os.getenv("HF_TOKEN")
 }
 
-
-def build_graph(provider, model_config):
-    from langchain_core.messages import SystemMessage, HumanMessage
-    #from langgraph.graph import StateGraph, ToolNode
-    from langchain_core.runnables import RunnableLambda
-    from some_module import vector_store  # Make sure this is defined/imported
-
-    # Step 1: Get LLM
+    # Get LLM
     def get_llm(provider: str, config: dict):
         if provider == "huggingface":
             from langchain_huggingface import HuggingFaceEndpoint
@@ -1056,7 +1072,30 @@ def build_graph(provider, model_config):
             )
         else:
             raise ValueError(f"Unsupported provider: {provider}")
-    
+
+
+    def assistant(state: dict):
+        return {"messages": [llm_with_tools.invoke(state["messages"])]}
+
+    def tools_condition(state: dict) -> str:
+        if "use tool" in state["messages"][-1].content.lower():
+            return "tools"
+        else:
+            return "END" 
+            
+
+from langgraph.graph import StateGraph
+
+# Build graph using AgentState as the shared schema
+def build_graph() -> StateGraph:
+    builder = StateGraph(AgentState)
+
+    from langchain_core.messages import SystemMessage, HumanMessage
+    #from langgraph.graph import StateGraph, ToolNode
+    from langchain_core.runnables import RunnableLambda
+    from some_module import vector_store  # Make sure this is defined/imported
+
+
     llm = get_llm(provider, model_config)
 
     # Step 2: Define tools
@@ -1071,80 +1110,52 @@ def build_graph(provider, model_config):
         wikidata_query
     ]
 
+    global tool_map
+    tool_map = {t.name: t for t in tools}
+
     # Step 3: Bind tools to LLM
     llm_with_tools = llm.bind_tools(tools)
 
     # Step 4: Build stateful graph logic
     sys_msg = SystemMessage(content="You are a helpful assistant.")
 
-    def retriever(state: dict):
-        user_query = state["messages"][0].content
-        similar_docs = vector_store.similarity_search(user_query)
-
-        if not similar_docs:
-            wiki_result = wiki_search.run(user_query)
-            return {
-                "messages": [
-                    sys_msg,
-                    state["messages"][0],
-                    HumanMessage(content=f"Using Wikipedia search:\n\n{wiki_result}")
-                ]
-            }
-        else:
-            return {
-                "messages": [
-                    sys_msg,
-                    state["messages"][0],
-                    HumanMessage(content=f"Reference:\n\n{similar_docs[0].page_content}")
-                ]
-            }
-
-    def assistant(state: dict):
-        return {"messages": [llm_with_tools.invoke(state["messages"])]}
-
-    def tools_condition(state: dict) -> str:
-        if "use tool" in state["messages"][-1].content.lower():
-            return "tools"
-        else:
-            return "END"
-
+    # --- Define nodes ---
+    retriever = RunnableLambda(lambda state: {
+        **state,
+        "retrieved_docs": vector_store.similarity_search(state["input"])
+    })
 
+    assistant = RunnableLambda(lambda state: {
+        **state,
+        "messages": state["messages"] + [SystemMessage(content="You are a helpful assistant.")]
+    })
 
-from langgraph.graph import StateGraph
+    call_llm = llm.bind_tools(tools)
 
-# Build graph using AgentState as the shared schema
-builder = StateGraph(AgentState)
+    # --- Build the graph ---
+    builder = StateGraph(AgentState)
 
-# Add nodes
-builder.add_node("retriever", retriever)
-builder.add_node("assistant", assistant)
-builder.add_node("call_llm", call_llm)
-builder.add_node("call_tool", tool_dispatcher)  # one name is enough
+    builder.add_node("retriever", retriever)
+    builder.add_node("assistant", assistant)
+    builder.add_node("call_llm", call_llm)
+    builder.add_node("call_tool", tool_dispatcher)
 
-# Entry point
-builder.set_entry_point("retriever")
-
-# Define the flow
-builder.add_edge("retriever", "assistant")
-builder.add_edge("assistant", "call_llm")
-
-# Conditional edge from LLM to tool or end
-builder.add_conditional_edges("call_llm", should_call_tool, {
-    "call_tool": "call_tool",
-    "end": None
-})
-
-# Loop back after tool execution
-builder.add_edge("call_tool", "call_llm")
+    builder.set_entry_point("retriever")
+    builder.add_edge("retriever", "assistant")
+    builder.add_edge("assistant", "call_llm")
+    builder.add_conditional_edges("call_llm", should_call_tool, {
+        "call_tool": "call_tool",
+        "end": None
+    })
+    builder.add_edge("call_tool", "call_llm")
 
-# Compile
-graph = builder.compile()
-return graph
+    graph = builder.compile()
+    return graph
 
 
 
 # call build_graph AFTER it’s defined
-agent = build_graph(provider, model_config)
+agent = graph
 
 # Now you can use the agent like this:
 result = agent.invoke({"messages": [HumanMessage(content=question)]})