Update app.py

app.py

@@ -137,57 +137,55 @@ def should_continue(state: AgentState) -> str:
     return "continue"
 
 
-
 def reasoning_node(state: AgentState) -> AgentState:
     import os
+    from langchain_google_genai import ChatGoogleGenerativeAI
     from langchain.schema import HumanMessage, AIMessage
-    from langchain.prompts import PromptTemplate
-    from langchain_community.llms import HuggingFaceHub
-    from langchain.chains import LLMChain
+    from langchain.prompts import ChatPromptTemplate
 
-    # 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")
+    # Explicitly load the Google 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.")
 
-    # Defensive: Ensure valid history
+    # Defensive: Ensure valid message history
     if not state["history"] or not isinstance(state["history"][-1], HumanMessage):
         state["history"].append(HumanMessage(content="Continue."))
 
-    # Create the LLM (Zephyr-7B)
-    llm = HuggingFaceHub(
-        repo_id="HuggingFaceH4/zephyr-7b-beta",
-        huggingfacehub_api_token=token,
-        model_kwargs={"temperature": 0.1, "max_new_tokens": 500}
-    )
-
-    # 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]"
+    # Initialize Gemini model
+    llm = ChatGoogleGenerativeAI(
+        model="gemini-1.5-flash",
+        temperature=0.1,
+        google_api_key=GOOGLE_API_KEY
     )
 
-    # Build the chain
-    chain = LLMChain(prompt=flat_prompt, llm=llm)
-
-    # Run the chain
+    # 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 | llm
+
+    # Invoke chain with inputs
     response = chain.invoke({
         "context": state["context"],
         "reasoning": state["reasoning"],
         "question": state["question"]
     })
 
-    content = response["text"]
-
-    # Parse response
+    content = response.content
     reasoning, action, action_input = parse_agent_response(content)
 
-    # Update state
+    # Update agent state
     state['history'].append(AIMessage(content=content))
     state['reasoning'] += f"\nStep {state['iterations']+1}: {reasoning}"
     state['iterations'] += 1
@@ -202,7 +200,6 @@ def reasoning_node(state: AgentState) -> AgentState:
 
     return state
 
-
     
 
 def tool_node(state: AgentState) -> AgentState: