Update agent.py

agent.py

@@ -1,14 +1,26 @@
 import os
+from dotenv import load_dotenv
 import gradio as gr
+
+
+# LangGraph & LangChain
+from langgraph.graph import START, StateGraph, MessagesState
+from langgraph.prebuilt import ToolNode, tools_condition
+from langchain_core.messages import SystemMessage, HumanMessage
+from langchain_core.tools import tool
+
+#infrence provider
 from langchain_huggingface import HuggingFaceEndpoint
+# Web search tool
 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
+
+# # 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
+
 
 # --- 1. LOAD API KEYS ---
 load_dotenv()
@@ -22,31 +34,153 @@ 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)
+# tools = [TavilySearchResults(max_results=3), PythonREPLTool()]
+# tool_node = ToolNode(tools)
+
+
+### TOOLS
+
+@tool
+def multiply(a: int, b: int) -> int:
+    """Multiply two numbers.
+    Args:
+        a: first int
+        b: second int
+    """
+    return a * b
+
+@tool
+def add(a: int, b: int) -> int:
+    """Add two numbers.
+    
+    Args:
+        a: first int
+        b: second int
+    """
+    return a + b
+
+@tool
+def subtract(a: int, b: int) -> int:
+    """Subtract two numbers.
+    
+    Args:
+        a: first int
+        b: second int
+    """
+    return a - b
+
+@tool
+def divide(a: int, b: int) -> int:
+    """Divide two numbers.
+    
+    Args:
+        a: first int
+        b: second int
+    """
+    if b == 0:
+        raise ValueError("Cannot divide by zero.")
+    return a / b
+
+@tool
+def modulus(a: int, b: int) -> int:
+    """Get the modulus of two numbers.
+    
+    Args:
+        a: first int
+        b: second int
+    """
+    return a % b
+
+
+
+@tool
+def web_search(query: str) -> str:
+    """Search Tavily for a query and return maximum 3 results.
+    
+    Args:
+        query: The search query."""
+    search_docs = TavilySearchResults(max_results=3).invoke(query=query)
+    formatted_search_docs = "\n\n---\n\n".join(
+        [
+            f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
+            for doc in search_docs
+        ])
+    return {"web_results": formatted_search_docs}
 
-repo_id = "meta-llama/Meta-Llama-3-8B-Instruct"
-llm = HuggingFaceEndpoint(
+# 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. 
+"""
+tools = [divide, add, multiply,subtract, web_search,modulus ]
+
+### LLM
+# 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)
+
+
+def build_graph():
+    """Builds and returns the LangGraph graph."""
+    #llm = ChatGroq(model="qwen-qwq-32b", temperature=0,api_key=groq_api_key)
+    repo_id = "meta-llama/Meta-Llama-3-8B-Instruct"
+    llm = HuggingFaceEndpoint(
     repo_id=repo_id,
     huggingfacehub_api_token=hf_token,
-    temperature=0.1,
+    temperature=0,
     max_new_tokens=1024,
 )
-llm_with_tools = llm.bind_tools(tools)
+    llm_with_tools = llm.bind_tools(tools)
+
+# Node
+    def assistant(state: MessagesState):
+        """Assistant node"""
+        return {"messages": [llm_with_tools.invoke([system_prompt] + state["messages"])]}
+
+
+    builder = StateGraph(MessagesState)
+    # Nodes
+    builder.add_node("assistant", assistant)
+    builder.add_node("tools", ToolNode(tools))
+
+    # Edges
+    builder.add_edge(START, "assistant")
+    builder.add_conditional_edges("assistant", tools_condition)
+    builder.add_edge("tools", "assistant")
+
+    #Compile graph
+    return builder.compile()
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
 
 
 # --- 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. 
-"""
+
 
 
 
@@ -115,4 +249,5 @@ iface = gr.Interface(
 
 # --- 8. LAUNCH THE APP ---
 iface.launch()
-# gr.launch()
+# gr.launch()
+"""