Update basic_agent.py

basic_agent.py

@@ -1,272 +1,189 @@
-"""
-A LangGraph-powered agent with at least 14 tools for question answering.
-"""
-import wikipedia
-import sympy
-import requests
-from datetime import datetime
-from dateutil import parser as date_parser
-from duckduckgo_search import DDGS
-import pytz
-import openai
-import wolframalpha
-import pandas as pd
-from bs4 import BeautifulSoup
-import httpx
-from langgraph.graph import StateGraph, END
-
-# 1. Wikipedia Search Tool
-def wikipedia_search(query):
-    try:
-        return wikipedia.summary(query, sentences=2)
-    except Exception as e:
-        return f"Wikipedia error: {e}"
-
-# 2. Math Evaluation Tool
-def math_eval(expr):
-    try:
-        return str(sympy.sympify(expr).evalf())
-    except Exception as e:
-        return f"Math error: {e}"
-
-# 3. DuckDuckGo Web Search Tool
-def ddg_search(query):
-    try:
-        results = list(DDGS().text(query, max_results=1))
-        return results[0]['body'] if results else 'No results.'
-    except Exception as e:
-        return f"DDG error: {e}"
-
-# 4. Current Date Tool
-def get_current_date():
-    return datetime.now().strftime('%Y-%m-%d')
-
-# 5. Current Time Tool
-def get_current_time():
-    return datetime.now().strftime('%H:%M:%S')
-
-# 6. Timezone Conversion Tool
-def convert_timezone(dt_str, from_tz, to_tz):
-    try:
-        dt = date_parser.parse(dt_str)
-        from_zone = pytz.timezone(from_tz)
-        to_zone = pytz.timezone(to_tz)
-        dt = from_zone.localize(dt)
-        return dt.astimezone(to_zone).strftime('%Y-%m-%d %H:%M:%S')
-    except Exception as e:
-        return f"Timezone error: {e}"
-
-# 7. OpenAI GPT-3.5 Completion Tool (requires OPENAI_API_KEY env var)
-def openai_completion(prompt):
-    try:
-        openai.api_key = os.getenv('OPENAI_API_KEY')
-        resp = openai.ChatCompletion.create(model="gpt-3.5-turbo", messages=[{"role": "user", "content": prompt}])
-        return resp.choices[0].message.content.strip()
-    except Exception as e:
-        return f"OpenAI error: {e}"
-
-# 8. WolframAlpha Short Answer Tool (requires WOLFRAMALPHA_APPID env var)
-def wolfram_query(query):
-    try:
-        appid = os.getenv('WOLFRAMALPHA_APPID')
-        client = wolframalpha.Client(appid)
-        res = client.query(query)
-        return next(res.results).text
-    except Exception as e:
-        return f"WolframAlpha error: {e}"
-
-# 9. Weather API Tool (using python-weather-api)
-def get_weather(city):
-    try:
-        import python_weather
-        import asyncio
-        async def getweather():
-            async with python_weather.Client(unit=python_weather.METRIC) as client:
-                weather = await client.get(city)
-                return f"{weather.current.temperature}°C, {weather.current.sky_text}"
-        return asyncio.run(getweather())
-    except Exception as e:
-        return f"Weather error: {e}"
-
-# 10. Pandas DataFrame Tool
-def describe_dataframe(csv_url):
-    try:
-        df = pd.read_csv(csv_url)
-        return str(df.describe())
-    except Exception as e:
-        return f"Pandas error: {e}"
-
-# 11. BeautifulSoup HTML Title Extractor
-def extract_title(url):
-    try:
-        resp = requests.get(url, timeout=10)
-        soup = BeautifulSoup(resp.text, 'lxml')
-        return soup.title.string.strip() if soup.title else 'No title found.'
-    except Exception as e:
-        return f"Soup error: {e}"
-
-# 12. HTTPX GET Tool
-def httpx_get(url):
-    try:
-        resp = httpx.get(url, timeout=10)
-        return resp.text[:500]
-    except Exception as e:
-        return f"HTTPX error: {e}"
-
-# 13. Currency Conversion Tool (using exchangerate.host)
-def convert_currency(amount, from_cur, to_cur):
-    try:
-        url = f"https://api.exchangerate.host/convert?from={from_cur}&to={to_cur}&amount={amount}"
-        resp = requests.get(url)
-        data = resp.json()
-        return f"{amount} {from_cur} = {data['result']} {to_cur}"
-    except Exception as e:
-        return f"Currency error: {e}"
-
-# 14. IP Geolocation Tool
-def ip_geolocate(ip):
-    try:
-        resp = requests.get(f"https://ipinfo.io/{ip}/json")
-        data = resp.json()
-        return f"{data.get('city', '?')}, {data.get('region', '?')}, {data.get('country', '?')}"
-    except Exception as e:
-        return f"IP error: {e}"
-
-# --- LangGraph Agent Setup ---
-from langgraph.agent import Tool, Agent
-
-tools = [
-    Tool("wikipedia_search", wikipedia_search, description="Search Wikipedia for a summary."),
-    Tool("math_eval", math_eval, description="Evaluate a math expression."),
-    Tool("ddg_search", ddg_search, description="DuckDuckGo web search."),
-    Tool("get_current_date", get_current_date, description="Get the current date."),
-    Tool("get_current_time", get_current_time, description="Get the current time."),
-    Tool("convert_timezone", convert_timezone, description="Convert time between timezones."),
-    Tool("openai_completion", openai_completion, description="Get a completion from OpenAI GPT-3.5."),
-    Tool("wolfram_query", wolfram_query, description="Query WolframAlpha for a short answer."),
-    Tool("get_weather", get_weather, description="Get current weather for a city."),
-    Tool("describe_dataframe", describe_dataframe, description="Describe a CSV file using pandas."),
-    Tool("extract_title", extract_title, description="Extract the title from a webpage."),
-    Tool("httpx_get", httpx_get, description="Fetch a webpage using HTTPX."),
-    Tool("convert_currency", convert_currency, description="Convert currency using exchangerate.host."),
-    Tool("ip_geolocate", ip_geolocate, description="Get geolocation info for an IP address."),
-]
-
-class LangGraphAgent:
-    def __init__(self):
-        self.agent = Agent(tools=tools)
-    def __call__(self, question: str) -> str:
-        """Use LangGraph agent to answer the question."""
-        try:
-            return self.agent.run(question)
-        except Exception as e:
-            return f"LangGraphAgent error: {e}"
-
-# Define a simple state for demonstration (can be extended for more complex workflows)
-class AgentState(dict):
-    pass
-
-def build_langgraph():
-    # Create a graph
-    graph = StateGraph(AgentState)
+"""LangGraph Agent"""
+import os
+from dotenv import load_dotenv
+from langgraph.graph import START, StateGraph, MessagesState, END
+from langgraph.prebuilt import tools_condition
+from langgraph.prebuilt import ToolNode
+from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain_groq import ChatGroq
+from langchain_huggingface import ChatHuggingFace, HuggingFaceEndpoint, HuggingFaceEmbeddings
+from langchain_community.tools.tavily_search import TavilySearchResults
+from langchain_community.document_loaders import WikipediaLoader
+from langchain_community.document_loaders import ArxivLoader
+from langchain_community.vectorstores import SupabaseVectorStore
+from langchain_core.messages import SystemMessage, HumanMessage, AIMessage
+from langchain_core.tools import tool
+from langchain.tools.retriever import create_retriever_tool
+from supabase.client import Client, create_client
+from typing import List, Dict, Any
+
+load_dotenv()
+
+# --- Tools ---
+@tool
+def multiply(a: int, b: int) -> int:
+    """Multiply two numbers.
+    Args:
+        a: first int
+        b: second int
+    """
+    return a * b
 
-    # Add each tool as a node
-    for tool in tools:
-        def make_tool_node(tool_func):
-            def node(state: AgentState):
-                question = state.get('question', '')
-                try:
-                    result = tool_func(question)
-                except Exception as e:
-                    result = f"Tool error: {e}"
-                state['result'] = result
-                return state
-            return node
-        graph.add_node(tool.name, make_tool_node(tool.func))
+@tool
+def add(a: int, b: int) -> int:
+    """Add two numbers.
 
-    # For demonstration, connect all nodes to END (in practice, you may want to chain or branch)
-    for tool in tools:
-        graph.add_edge(tool.name, END)
+    Args:
+        a: first int
+        b: second int
+    """
+    return a + b
 
-    # Set entry point (first tool)
-    graph.set_entry_point(tools[0].name)
+@tool
+def subtract(a: int, b: int) -> int:
+    """Subtract two numbers.
 
-    return graph
+    Args:
+        a: first int
+        b: second int
+    """
+    return a - b
 
-# --- GAIA-style Multi-step Agent ---
-import re
+@tool
+def divide(a: int, b: int) -> int:
+    """Divide two numbers.
 
-class GAIAAgent:
+    Args:
+        a: first int
+        b: second int
     """
-    A simple GAIA-style agent that can plan and chain multiple tools for multi-step reasoning.
-    This is a minimal demonstration and can be extended for more advanced planning.
-    """
-    def __init__(self):
-        self.tool_map = {t.name: t.func for t in tools}
-
-    def plan(self, question: str):
-        """
-        Very simple planner: looks for keywords to select tools and chain them.
-        Returns a list of (tool_name, tool_input) tuples.
-        """
-        steps = []
-        q = question.lower()
-        # Example: If question asks for weather in a city and then convert time
-        if 'weather' in q and 'time' in q and 'convert' in q:
-            # e.g., "What is the weather in Paris and convert the time to Tokyo timezone?"
-            city = re.findall(r'weather in ([a-zA-Z ]+)', q)
-            city = city[0].strip() if city else 'London'
-            steps.append(('get_weather', city))
-            # Assume user wants to convert current time from city to Tokyo
-            steps.append(('convert_timezone', [datetime.now().strftime('%Y-%m-%d %H:%M:%S'), 'Europe/Paris', 'Asia/Tokyo']))
-        elif 'weather' in q:
-            city = re.findall(r'weather in ([a-zA-Z ]+)', q)
-            city = city[0].strip() if city else 'London'
-            steps.append(('get_weather', city))
-        elif 'currency' in q or 'convert' in q and 'usd' in q and 'eur' in q:
-            amount = re.findall(r'(\d+)', q)
-            amount = amount[0] if amount else '1'
-            steps.append(('convert_currency', [amount, 'USD', 'EUR']))
-        elif 'wikipedia' in q or 'who is' in q or 'what is' in q:
-            topic = re.findall(r'(?:wikipedia|who is|what is) ([a-zA-Z0-9 ]+)', q)
-            topic = topic[0].strip() if topic else question
-            steps.append(('wikipedia_search', topic))
-        elif 'math' in q or any(op in q for op in ['+', '-', '*', '/']):
-            expr = re.findall(r'([\d\s\+\-\*/\.]+)', q)
-            expr = expr[0] if expr else question
-            steps.append(('math_eval', expr))
-        else:
-            # Default: try DuckDuckGo search
-            steps.append(('ddg_search', question))
-        return steps
+    if b == 0:
+        raise ValueError("Cannot divide by zero.")
+    return a / b
 
-    def __call__(self, question: str) -> str:
-        """
-        Execute the planned steps, chaining outputs if needed.
-        """
-        steps = self.plan(question)
-        last_output = None
-        for tool_name, tool_input in steps:
-            func = self.tool_map.get(tool_name)
-            if not func:
-                last_output = f"Tool {tool_name} not found."
-                break
-            # If previous output is needed as input, use it
-            if isinstance(tool_input, list):
-                # Replace any placeholder with last_output
-                tool_input = [last_output if x == '__PREV__' else x for x in tool_input]
-                try:
-                    last_output = func(*tool_input)
-                except Exception as e:
-                    last_output = f"Error in {tool_name}: {e}"
-            else:
-                try:
-                    last_output = func(tool_input)
-                except Exception as e:
-                    last_output = f"Error in {tool_name}: {e}"
-        return last_output
+@tool
+def modulus(a: int, b: int) -> int:
+    """Get the modulus of two numbers.
 
-# Example usage:
-# graph = build_langgraph()
-# result = graph.run({'question': 'What is the capital of France?'})
-# print(result['result'])
+    Args:
+        a: first int
+        b: second int
+    """
+    return a % b
+
+@tool
+def wiki_search(query: str) -> str:
+    """Search Wikipedia for a query and return maximum 2 results.
+
+    Args:
+        query: The search query."""
+    search_docs = WikipediaLoader(query=query, load_max_docs=2).load()
+    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 {"wiki_results": formatted_search_docs}
+
+@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}
+
+@tool
+def arxiv_search(query: str) -> str:
+    """Search Arxiv for a query and return maximum 3 result.
+
+    Args:
+        query: The search query."""
+    search_docs = ArxivLoader(query=query, load_max_docs=3).load()
+    formatted_search_docs = "\n\n---\n\n".join(
+        [
+            f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content[:1000]}\n</Document>'
+            for doc in search_docs
+        ])
+    return {"arxiv_results": formatted_search_docs}
+
+
+# --- Prompt & Retriever Setup ---
+with open("system_prompt.txt", "r", encoding="utf-8") as f:
+    system_prompt = f.read()
+sys_msg = SystemMessage(content=system_prompt)
+
+embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
+supabase: Client = create_client(
+    os.environ.get("SUPABASE_URL"),
+    os.environ.get("SUPABASE_SERVICE_KEY"))
+vector_store = SupabaseVectorStore(
+    client=supabase,
+    embedding=embeddings,
+    table_name="documents",
+    query_name="match_documents_langchain",
+)
+retriever = vector_store.as_retriever()  # Access the retriever directly
+
+# --- Graph Definition ---
+def build_graph(provider: str = "google"):
+    """Build the graph"""
+    if provider == "google":
+        llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash", temperature=0)
+    elif provider == "groq":
+        llm = ChatGroq(model="qwen-qwq-32b", temperature=0)
+    elif provider == "huggingface":
+        llm = ChatHuggingFace(
+            llm=HuggingFaceEndpoint(
+                url="https://api-inference.huggingface.co/models/Meta-DeepLearning/llama-2-7b-chat-hf",
+                temperature=0,
+            ),
+        )
+    else:
+        raise ValueError("Invalid provider. Choose 'google', 'groq' or 'huggingface'.")
+
+    # Bind tools to LLM
+    llm_with_tools = llm.bind_tools(tools)
+
+    # Define nodes
+    def retrieval_node(state: MessagesState):
+        """Retrieves relevant documents."""
+        query = state["messages"][-1].content # Get latest message
+        docs = retriever.get_relevant_documents(query)
+        context = "\n\n".join([d.page_content for d in docs])  # Concatenate document content
+        #Append context to the messages so that it can be send to LLM
+        return {"messages": [HumanMessage(content=f"Context:\n{context}\nOriginal question: {query}")]}
+
+    def llm_node(state: MessagesState):
+        """Invokes the LLM to answer the question."""
+        return {"messages": [llm_with_tools.invoke(state)]}
+
+    # Add tools node
+    tool_node = ToolNode(tools)
+    # Define graph
+
+    builder = StateGraph(MessagesState)
+    builder.add_node("retrieval", retrieval_node)
+    builder.add_node("llm", llm_node)
+    builder.add_node("tools", tool_node)
+
+    # Graph structure
+    builder.set_entry_point("retrieval")
+    builder.add_edge("retrieval", "llm")
+
+    #Conditional Edges
+    builder.add_conditional_edges(
+        "llm",
+        tools_condition,
+    )
+    builder.add_edge("tools", "llm")
+    builder.add_edge("llm", END)
+
+    # Compile
+    return builder.compile()