Update app.py

app.py

@@ -28,6 +28,9 @@ from langchain_core.messages import HumanMessage, AIMessage, ToolMessage
 from langchain_core.tools import tool
 from typing import Dict, List, TypedDict, Annotated
 import operator
+from langchain_community.llms import HuggingFaceHub
+from langchain_community.chat_models import ChatHuggingFace
+
 
 # ====== Tool Definitions ======
 @tool
@@ -125,43 +128,44 @@ def should_continue(state: AgentState) -> str:
 
 def reasoning_node(state: AgentState) -> AgentState:
     """Agent reasoning and tool selection"""
-    from langchain_community.chat_models import ChatHuggingFace
-    from langchain.schema import SystemMessage
+    # Get Hugging Face API token from environment
+    token = os.environ.get("HF_API_TOKEN")
+    if not token:
+        raise ValueError("Hugging Face API token not found in environment variables")
+    
+    # Create the underlying LLM model
+    llm = HuggingFaceHub(
+        repo_id="HuggingFaceH4/zephyr-7b-beta",
+        huggingfacehub_api_token=token,
+        model_kwargs={
+            "temperature": 0.1,
+            "max_new_tokens": 500
+        }
+    )
+    
+    # Wrap the LLM in ChatHuggingFace
+    chat_model = ChatHuggingFace(llm=llm)
     
     # Build prompt
     prompt = ChatPromptTemplate.from_messages([
-        SystemMessage(content=(
-            "You are an intelligent AI assistant. Follow this process:\n"
-            "1. Analyze the question: {question}\n"
-            "2. Review context: {context}\n"
-            "3. Reasoning Steps:\n{reasoning}\n"
-            "4. Select ONE tool to use next OR provide FINAL ANSWER\n\n"
-            "Available Tools:\n"
-            "- duckduckgo_search: For current information\n"
-            "- wikipedia_search: For factual knowledge\n"
-            "- arxiv_search: For academic topics\n"
-            "- document_qa: For questions about documents\n"
-            "- python_execution: For calculations/code\n\n"
-            "Response Format:\n"
-            "Reasoning: [Your step-by-step analysis]\n"
-            "Action: [Tool name OR 'Final Answer']\n"
-            "Action Input: [Tool parameters OR final response]"
+        ("system", (
+            "You're an expert problem solver. Analyze the question, select the best tool, "
+            "and provide reasoning. Available tools: duckduckgo_search, wikipedia_search, "
+            "arxiv_search, document_qa, python_execution.\n\n"
+            "Current Context:\n{context}\n\n"
+            "Reasoning Steps:\n{reasoning}\n\n"
+            "Response Format:\nReasoning: [Your analysis]\nAction: [Tool name OR 'Final Answer']\n"
+            "Action Input: [Input for tool OR final response]"
         )),
         *state['history']
     ])
     
-    # Initialize model
-    llm = ChatHuggingFace(
-        model_name="HuggingFaceH4/zephyr-7b-beta",
-        temperature=0.3
-    )
-    
-    # Generate response
-    response = llm.invoke(prompt.format_messages(
-        question=state['question'],
-        context=state['context'],
-        reasoning=state['reasoning']
-    ))
+    chain = prompt | chat_model
+    response = chain.invoke({
+        "context": state['context'],
+        "reasoning": state['reasoning'],
+        "question": state['question']
+    })
     
     # Parse response
     content = response.content
@@ -172,17 +176,16 @@ def reasoning_node(state: AgentState) -> AgentState:
     state['reasoning'] += f"\nStep {state['iterations']+1}: {reasoning}"
     
     if "final answer" in action.lower():
-        state['history'].append(AIMessage(
-            content=f"FINAL ANSWER: {action_input}"
-        ))
+        state['history'].append(AIMessage(content=f"FINAL ANSWER: {action_input}"))
     else:
         state['history'].append({
-            "role": "action_request",
             "tool": action,
-            "input": action_input
+            "input": action_input,
+            "role": "action_request"
         })
     
     return state
+    
 
 def tool_node(state: AgentState) -> AgentState:
     """Execute selected tool and update state"""