Update agent.py

agent.py

@@ -0,0 +1,117 @@
+import os
+import gradio as gr
+from langchain_huggingface import HuggingFaceEndpoint
+from langchain_community.tools.tavily_search import TavilySearchResults
+# NEW IMPORT
+from langchain_experimental.tools import PythonREPLTool
+from langchain_core.messages import BaseMessage, HumanMessage
+from langgraph.graph import StateGraph, END
+from langgraph.prebuilt import ToolNode
+from typing import TypedDict, Annotated, List
+from dotenv import load_dotenv
+
+# --- 1. LOAD API KEYS ---
+load_dotenv()
+hf_token = os.getenv("HF_TOKEN")
+tavily_api_key = os.getenv("TAVILY_API_KEY")
+
+if not hf_token or not tavily_api_key:
+    raise ValueError("Hugging Face Token or Tavily API Key is not set in the environment variables.")
+os.environ["TAVILY_API_KEY"] = tavily_api_key
+
+
+# --- 2. DEFINE TOOLS and INITIALIZE LLM ---
+# UPDATED TOOLS LIST
+tools = [TavilySearchResults(max_results=3), PythonREPLTool()]
+tool_node = ToolNode(tools)
+
+repo_id = "meta-llama/Meta-Llama-3-8B-Instruct"
+llm = HuggingFaceEndpoint(
+    repo_id=repo_id,
+    huggingfacehub_api_token=hf_token,
+    temperature=0.1,
+    max_new_tokens=1024,
+)
+llm_with_tools = llm.bind_tools(tools)
+
+
+# --- 3. DEFINE THE AGENT'S STATE ---
+class AgentState(TypedDict):
+    messages: Annotated[List[BaseMessage], lambda x, y: x + y]
+
+# SYSTEM PROMPT
+system_prompt = """
+You are a helpful assistant tasked with answering questions using a set of tools. 
+Now, I will ask you a question. Report your thoughts, and finish your answer with the following template: 
+FINAL ANSWER: [YOUR FINAL ANSWER]. 
+YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise. If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string.
+Your answer should only start with "FINAL ANSWER: ", then follows with the answer. 
+"""
+
+
+
+# --- 4. DEFINE THE NODES OF THE GRAPH ---
+def agent_node(state):
+    response = llm_with_tools.invoke(state["messages"])
+    return {"messages": [response]}
+
+
+# --- 5. DEFINE THE EDGES OF THE GRAPH ---
+def should_continue(state):
+    last_message = state["messages"][-1]
+    if last_message.tool_calls:
+        return "tools"
+    return END
+
+
+# --- 6. ASSEMBLE THE GRAPH ---
+workflow = StateGraph(AgentState)
+
+workflow.add_node("agent", agent_node)
+workflow.add_node("tools", tool_node)
+
+workflow.set_entry_point("agent")
+
+workflow.add_conditional_edges(
+    "agent",
+    should_continue,
+    {
+        "tools": "tools",
+        "end": END,
+    },
+)
+
+workflow.add_edge("tools", "agent")
+
+app = workflow.compile()
+
+
+# --- 7. CREATE THE USER INTERFACE (UI) ---
+def run_agent(query: str):
+    try:
+        inputs = {"messages": [HumanMessage(content=query)]}
+        final_response = None
+        # Using stream to get final output, can be slow for complex tasks
+        for s in app.stream(inputs, {"recursion_limit": 10}):
+            if "agent" in s:
+                final_response = s["agent"]["messages"][-1].content
+        return final_response if final_response else "Agent did not produce a final answer."
+
+    except Exception as e:
+        return f"An error occurred: {e}"
+
+# iface = gr.Interface(
+#     fn=run_agent,
+#     inputs=gr.Textbox(lines=2, placeholder="Ask the agent anything..."),
+#     outputs="markdown",
+#     title="GAIA Agent v0.3 (LangGraph + Code Interpreter)",
+#     description="This agent can use web search and a Python code interpreter.",
+#     examples=[
+#         ["What is the square root of the number of states in the USA?"],
+#         ["What is the total number of letters in the names of the first three planets in our solar system?"]
+#     ],
+# )
+
+
+# # --- 8. LAUNCH THE APP ---
+# iface.launch()