Update app.py

app.py

@@ -141,11 +141,11 @@ def should_continue(state: AgentState) -> str:
 def reasoning_node(state: AgentState) -> AgentState:
     import os
     from langchain.schema import HumanMessage, AIMessage
-    from langchain_community.chat_models import ChatHuggingFace
-    from langchain.prompts import ChatPromptTemplate
+    from langchain.prompts import PromptTemplate
     from langchain_community.llms import HuggingFaceHub
+    from langchain.chains import LLMChain
 
-    # Ensure token is available
+    # Ensure HF token is available
     token = os.environ.get("HF_TOKEN")
     if not token:
         raise ValueError("Hugging Face API token not found in environment variables")
@@ -154,39 +154,40 @@ def reasoning_node(state: AgentState) -> AgentState:
     if not state["history"] or not isinstance(state["history"][-1], HumanMessage):
         state["history"].append(HumanMessage(content="Continue."))
 
-    # Create the LLM
+    # Create the LLM (Zephyr-7B)
     llm = HuggingFaceHub(
         repo_id="HuggingFaceH4/zephyr-7b-beta",
-        huggingfacehub_api_token=os.environ["HF_TOKEN"],
+        huggingfacehub_api_token=token,
         model_kwargs={"temperature": 0.1, "max_new_tokens": 500}
     )
 
-    chat_model = ChatHuggingFace(llm=llm)
-
-    # Build prompt
-    prompt = ChatPromptTemplate.from_messages([
-        ("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']
-    ])
-
-    chain = prompt | chat_model
+    # Flattened prompt (not chat-based)
+    flat_prompt = PromptTemplate.from_template(
+        "You're an expert problem solver. Analyze the question, select the best tool, "
+        "and provide reasoning.\n\n"
+        "Context:\n{context}\n\n"
+        "Reasoning Steps:\n{reasoning}\n\n"
+        "Question:\n{question}\n\n"
+        "Response Format:\nReasoning: [Your analysis]\nAction: [Tool name OR 'Final Answer']\n"
+        "Action Input: [Input for tool OR final response]"
+    )
+
+    # Build the chain
+    chain = LLMChain(prompt=flat_prompt, llm=llm)
+
+    # Run the chain
     response = chain.invoke({
-        "context": state['context'],
-        "reasoning": state['reasoning'],
-        "question": state['question']
+        "context": state["context"],
+        "reasoning": state["reasoning"],
+        "question": state["question"]
     })
 
-    content = response.content
+    content = response["text"]
+
+    # Parse response
     reasoning, action, action_input = parse_agent_response(content)
 
+    # Update state
     state['history'].append(AIMessage(content=content))
     state['reasoning'] += f"\nStep {state['iterations']+1}: {reasoning}"
     state['iterations'] += 1
@@ -194,7 +195,6 @@ def reasoning_node(state: AgentState) -> AgentState:
     if "final answer" in action.lower():
         state['history'].append(AIMessage(content=f"FINAL ANSWER: {action_input}"))
     else:
-        # Save action tool call in context instead of history
         state['context']['current_tool'] = {
             "tool": action,
             "input": action_input
@@ -202,6 +202,7 @@ def reasoning_node(state: AgentState) -> AgentState:
 
     return state
 
+
     
 
 def tool_node(state: AgentState) -> AgentState: