Update app.py

app.py

@@ -152,12 +152,12 @@ import os
 import time
 import json
 from typing import TypedDict, List, Union, Any, Dict
-from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain_huggingface import ChatHuggingFace
+from langchain_huggingface.llms import HuggingFaceEndpoint
 from langchain.schema import HumanMessage, AIMessage, SystemMessage
 from langchain.prompts import ChatPromptTemplate
 from langgraph.graph import StateGraph, END
-from google.api_core.exceptions import ResourceExhausted
-from langchain.tools import Tool # Import Tool for consistent tool definitions
+from langchain.tools import Tool
 
 # Assume these tools are defined elsewhere and imported
 # Placeholder for your actual tool implementations
@@ -248,7 +248,7 @@ def should_continue(state: AgentState) -> str:
     print(f"DEBUG: Entering should_continue. Current context: {state.get('context', {})}")
     
     # End if agent has produced a final answer
-    if state.get("final_answer"):
+    if state.get("final_answer") is not None: # Check for None explicitly
         print("DEBUG: should_continue -> END (Final Answer set in state)")
         return "end"
 
@@ -270,29 +270,31 @@ def reasoning_node(state: AgentState) -> AgentState:
     print(f"DEBUG: Entering reasoning_node. Iteration: {state['iterations']}")
     print(f"DEBUG: Current history length: {len(state.get('history', []))}")
 
-    # Load API key
-    GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
-    if not GOOGLE_API_KEY:
-        raise ValueError("GOOGLE_API_KEY not set in environment variables.")
+    # --- CHANGE: Use HF_TOKEN environment variable ---
+    HF_TOKEN = os.getenv("HF_TOKEN")
+    if not HF_TOKEN:
+        raise ValueError("HF_TOKEN not set in environment variables.")
+    # --- END CHANGE ---
 
-    # Initialize/update state fields
     state.setdefault("context", {})
     state.setdefault("reasoning", "")
     state.setdefault("iterations", 0)
     state.setdefault("current_task", "Understand the question and plan the next step.")
     state.setdefault("current_thoughts", "")
     
-    # Clear any old pending action from context before generating a new one
     state["context"].pop("pending_action", None) 
 
-    # Create Gemini model wrapper
-    llm = ChatGoogleGenerativeAI(
-        model="gemini-1.5-flash",
-        temperature=0.1,
-        google_api_key=GOOGLE_API_KEY
+    model_id = "meta-llama/Llama-2-7b-chat-hf" 
+    
+    llm = ChatHuggingFace(
+        llm=HuggingFaceEndpoint(
+            repo_id=model_id,
+            max_new_tokens=512, 
+            temperature=0.1,
+            huggingfacehub_api_token=HF_TOKEN, # --- CHANGE: Pass HF_TOKEN here ---
+        )
     )
 
-    # Dynamically generate tool descriptions for the prompt
     tool_descriptions = "\n".join([
         f"- **{t.name}**: {t.description}" for t in state.get("tools", [])
     ])
@@ -329,34 +331,30 @@ def reasoning_node(state: AgentState) -> AgentState:
 
     prompt = ChatPromptTemplate.from_messages([
         SystemMessage(content=system_prompt),
-        *state["history"] # Include full history for conversational context
+        *state["history"] 
     ])
 
     chain = prompt | llm
 
-    def call_with_retry(inputs, retries=3, delay=60):
+    def call_with_retry(inputs, retries=3, delay=30): 
         for attempt in range(retries):
             try:
                 response = chain.invoke(inputs)
-                # Attempt to parse immediately to catch bad JSON before returning
-                json.loads(response.content) # Validate JSON structure
+                json.loads(response.content) 
                 return response
-            except ResourceExhausted as e:
-                print(f"[Retry {attempt+1}/{retries}] Gemini rate limit hit. Waiting {delay}s...")
-                time.sleep(delay)
             except json.JSONDecodeError as e:
                 print(f"[Retry {attempt+1}/{retries}] LLM returned invalid JSON. Retrying...")
                 print(f"Invalid JSON content: {response.content[:200]}...")
-                time.sleep(5) # Shorter delay for parsing errors
-            except Exception as e:
-                print(f"[Retry {attempt+1}/{retries}] An unexpected error occurred during LLM call: {e}. Retrying...")
+                time.sleep(5) 
+            except Exception as e: 
+                print(f"[Retry {attempt+1}/{retries}] An unexpected error occurred during LLM call: {e}. Waiting {delay}s...")
                 time.sleep(delay)
-        raise RuntimeError("Failed after multiple retries due to Gemini quota limit or invalid JSON.")
+        raise RuntimeError("Failed after multiple retries due to Hugging Face API issues or invalid JSON.")
 
     response = call_with_retry({
         "context": state["context"],
         "reasoning": state["reasoning"],
-        "question": state["question"], # Redundant as it's in history, but keeps prompt consistent
+        "question": state["question"],
         "current_task": state["current_task"],
         "current_thoughts": state["current_thoughts"]
     })
@@ -367,22 +365,18 @@ def reasoning_node(state: AgentState) -> AgentState:
     print(f"DEBUG: LLM Raw Response Content: {content[:200]}...")
     print(f"DEBUG: Parsed Action: '{action}', Action Input: '{action_input[:100]}...'")
 
-    # Update state
-    state["history"].append(AIMessage(content=content)) # Store the raw LLM response
+    state["history"].append(AIMessage(content=content)) 
     state["reasoning"] += f"\nStep {state['iterations'] + 1}: {reasoning}"
     state["iterations"] += 1
-    state["current_thoughts"] = reasoning # Update current thoughts for next iteration
+    state["current_thoughts"] = reasoning 
 
     if "final answer" in action.lower():
-        state["final_answer"] = action_input # Set final answer directly in state
-        # The should_continue check will handle ending the graph based on final_answer presence
+        state["final_answer"] = action_input 
     else:
-        # Store the action request in context, not in history
         state["context"]["pending_action"] = {
             "tool": action,
             "input": action_input
         }
-        # Add a message to history to indicate the agent's intent for the LLM
         state["history"].append(AIMessage(content=f"Agent decided to use tool: {action} with input: {action_input}"))
 
 
@@ -396,12 +390,9 @@ def tool_node(state: AgentState) -> AgentState:
     """
     print(f"DEBUG: Entering tool_node. Iteration: {state['iterations']}")
 
-    # Get the pending action from context
     tool_call_dict = state["context"].pop("pending_action", None)
 
     if not tool_call_dict:
-        # This case should ideally not be reached if should_continue is robust,
-        # but provides a fallback.
         error_message = "[Tool Error] No pending_action found in context. This indicates an issue with graph flow."
         print(f"ERROR: {error_message}")
         state["history"].append(AIMessage(content=error_message))
@@ -410,34 +401,29 @@ def tool_node(state: AgentState) -> AgentState:
     tool_name = tool_call_dict.get("tool")
     tool_input = tool_call_dict.get("input")
 
-    # Defensive check for empty tool name or input (still needed as LLM might generate empty strings)
     if not tool_name or tool_input is None:
         error_message = f"[Tool Error] Invalid action request from LLM: Tool name '{tool_name}' or input '{tool_input}' was empty. LLM needs to provide valid 'Action' and 'Action Input'."
         print(f"ERROR: {error_message}")
         state["history"].append(AIMessage(content=error_message))
-        # Clear any problematic pending action
         state["context"].pop("pending_action", None)
         return state
 
-    # Look up and invoke the tool from the state's tool list
     available_tools = state.get("tools", [])
     tool_fn = next((t for t in available_tools if t.name == tool_name), None)
 
     if tool_fn is None:
-        # Fallback for unrecognized tool - feedback to LLM
         tool_output = f"[Tool Error] Tool '{tool_name}' not found or not available. Please choose from: {', '.join([t.name for t in available_tools])}"
         print(f"ERROR: {tool_output}")
     else:
         try:
             print(f"DEBUG: Invoking tool '{tool_name}' with input: '{tool_input[:100]}...'")
             tool_output = tool_fn.run(tool_input)
-            if not tool_output and tool_output is not False: # Ensure 'False' is not treated as empty
+            if not tool_output and tool_output is not False: 
                 tool_output = f"[{tool_name} output] No specific result found for '{tool_input}'. The tool might have returned an empty response."
         except Exception as e:
             tool_output = f"[Tool Error] An error occurred while running '{tool_name}': {str(e)}"
             print(f"ERROR: {tool_output}")
 
-    # Add tool output to history as an AIMessage for the LLM to process next
     state["history"].append(AIMessage(content=f"[{tool_name} output]\n{tool_output}"))
     
     print(f"DEBUG: Exiting tool_node. Tool output added to history. New history length: {len(state['history'])}")
@@ -445,17 +431,14 @@ def tool_node(state: AgentState) -> AgentState:
 
 
 # ====== Agent Graph ======
-def create_agent_workflow(tools: List[Tool]): # tools are passed in now
+def create_agent_workflow(tools: List[Tool]):
     workflow = StateGraph(AgentState)
     
-    # Define nodes
     workflow.add_node("reason", reasoning_node)
     workflow.add_node("action", tool_node)
     
-    # Set entry point
     workflow.set_entry_point("reason")
     
-    # Define edges
     workflow.add_conditional_edges(
         "reason",
         should_continue,
@@ -500,8 +483,6 @@ class BasicAgent:
             "tools": self.tools
         }
 
-        # The invoke method will now return the final state, or raise an error if it hits a dead end
-        # LangGraph runs are synchronous by default here.
         final_state = self.workflow.invoke(state)
 
         if final_state.get("final_answer") is not None:
@@ -509,8 +490,6 @@ class BasicAgent:
             print(f"--- Agent returning FINAL ANSWER: {answer} ---")
             return answer
         else:
-            # This should ideally not happen if the agent is designed to always provide a final answer
-            # or a specific "cannot answer" message.
             print(f"--- ERROR: Agent finished without setting 'final_answer' for question: {question} ---")
             raise ValueError("Agent finished without providing a final answer.")