Update agent.py

agent.py

@@ -1,114 +1,124 @@
-"""LangGraph Agent with OpenAI"""
+"""LangGraph Agent – versione GPT-4.1 / Hugging Face Spaces"""
 import os
+from dotenv import load_dotenv
 from langgraph.graph import START, StateGraph, MessagesState
-from langgraph.prebuilt import tools_condition, ToolNode
-from langchain_openai import ChatOpenAI
-from langchain_community.document_loaders import WikipediaLoader
-from langchain_community.document_loaders import ArxivLoader
-from langchain_core.messages import SystemMessage, HumanMessage, AIMessage
+from langgraph.prebuilt import tools_condition
+from langgraph.prebuilt import ToolNode
+
+# LLM providers
+from langchain_openai import ChatOpenAI            # NEW (GPT-4.1)
+from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain_groq import ChatGroq
+from langchain_huggingface import (
+    ChatHuggingFace,
+    HuggingFaceEndpoint,
+    HuggingFaceEmbeddings,
+)
+
+# Tools & loaders
+from langchain_community.tools.tavily_search import TavilySearchResults
+from langchain_community.document_loaders import WikipediaLoader, ArxivLoader
+from langchain_community.vectorstores import SupabaseVectorStore
+from langchain_core.messages import SystemMessage, HumanMessage
 from langchain_core.tools import tool
+from langchain.tools.retriever import create_retriever_tool
+from supabase.client import Client, create_client
 
-# Tools definition
+# --------------------------------------------------------------------------- #
+#  Carica variabili d'ambiente (.env eventuale + secrets di HF Spaces)        #
+# --------------------------------------------------------------------------- #
+load_dotenv()  # nei Spaces le secrets sono già in os.environ
+
+# --------------------------------------------------------------------------- #
+#  TOOL di esempio (aritmetica)                                               #
+# --------------------------------------------------------------------------- #
 @tool
-def multiply(a: int, b: int) -> int:
-    """Multiply two numbers.
-    
-    Args:
-        a: first int
-        b: second int
-    """
-    return a * b
+def multiply(a: int, b: int) -> 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
+def add(a: int, b: int) -> 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
+def subtract(a: int, b: int) -> int: return a - b
 
 @tool
 def divide(a: int, b: int) -> float:
-    """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
+def modulus(a: int, b: int) -> int: return a % b
 
+# --------------------------------------------------------------------------- #
+#  TOOL: Wikipedia                                                             #
+# --------------------------------------------------------------------------- #
 @tool
 def wiki_search(query: str) -> str:
-    """Search Wikipedia for a query and return maximum 2 results.
-    
-    Args:
-        query: The search query.
-    """
-    try:
-        search_docs = WikipediaLoader(query=query, load_max_docs=2).load()
-        if not search_docs:
-            return f"No Wikipedia results found for: {query}"
-        
-        formatted_search_docs = "\n\n---\n\n".join(
-            [
-                f'Source: {doc.metadata.get("source", "Wikipedia")}\nContent: {doc.page_content[:2000]}...'
-                for doc in search_docs
-            ])
-        return formatted_search_docs
-    except Exception as e:
-        return f"Error searching Wikipedia: {str(e)}"
-
+    """Search Wikipedia (max 2 docs) and return formatted result."""
+    docs = WikipediaLoader(query=query, load_max_docs=2).load()
+    return "\n\n---\n\n".join(
+        f'<Document source="{d.metadata["source"]}" page="{d.metadata.get("page","")}"/>\n'
+        f"{d.page_content}\n</Document>"
+        for d in docs
+    )
 
+# --------------------------------------------------------------------------- #
+#  TOOL: Tavily web search                                                     #
+# --------------------------------------------------------------------------- #
+@tool
+def web_search(query: str) -> str:
+    """Search Tavily (max 3 docs) and return formatted result."""
+    docs = TavilySearchResults(max_results=3).invoke(query=query)
+    return "\n\n---\n\n".join(
+        f'<Document source="{d.metadata["source"]}" page="{d.metadata.get("page","")}"/>\n'
+        f"{d.page_content}\n</Document>"
+        for d in docs
+    )
 
+# --------------------------------------------------------------------------- #
+#  TOOL: ArXiv                                                                 #
+# --------------------------------------------------------------------------- #
 @tool
 def arxiv_search(query: str) -> str:
-    """Search Arxiv for a query and return maximum 3 results.
-    
-    Args:
-        query: The search query.
-    """
-    try:
-        search_docs = ArxivLoader(query=query, load_max_docs=3).load()
-        if not search_docs:
-            return f"No Arxiv results found for: {query}"
-            
-        formatted_search_docs = "\n\n---\n\n".join(
-            [
-                f'Title: {doc.metadata.get("Title", "Unknown")}\nAuthors: {doc.metadata.get("Authors", "Unknown")}\nContent: {doc.page_content[:1500]}...'
-                for doc in search_docs
-            ])
-        return formatted_search_docs
-    except Exception as e:
-        return f"Error searching Arxiv: {str(e)}"
-
-# System prompt
-system_prompt = """You are a general AI assistant. I will ask you a question. Do not report your thoughts or comments, and finish your answer with the following template: [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."""
-
-# Tools list
+    """Search ArXiv (max 3 docs) and return formatted snippet."""
+    docs = ArxivLoader(query=query, load_max_docs=3).load()
+    return "\n\n---\n\n".join(
+        f'<Document source="{d.metadata["source"]}" page="{d.metadata.get("page","")}"/>\n'
+        f"{d.page_content[:1000]}\n</Document>"
+        for d in docs
+    )
+
+# --------------------------------------------------------------------------- #
+#  System prompt                                                               #
+# --------------------------------------------------------------------------- #
+with open("system_prompt.txt", "r", encoding="utf-8") as f:
+    system_prompt = f.read()
+sys_msg = SystemMessage(content=system_prompt)
+
+# --------------------------------------------------------------------------- #
+#  Vector store per il retriever                                               #
+# --------------------------------------------------------------------------- #
+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",
+)
+question_search_tool = create_retriever_tool(
+    retriever=vector_store.as_retriever(),
+    name="Question Search",
+    description="A tool to retrieve similar questions from a vector store.",
+)
+
+# --------------------------------------------------------------------------- #
+#  Registrazione tool list                                                    #
+# --------------------------------------------------------------------------- #
 tools = [
     multiply,
     add,
@@ -116,125 +126,93 @@ tools = [
     divide,
     modulus,
     wiki_search,
+    web_search,
     arxiv_search,
+    question_search_tool,
 ]
 
-class LangGraphAgent:
-    """LangGraph Agent with OpenAI that can be used in HuggingFace Space evaluation"""
-    
-    def __init__(self):
-        """Initialize the agent with OpenAI LLM and tools"""
-        print("Initializing LangGraphAgent...")
-        
-        # Get API key from environment
-        self.api_key = os.environ.get("OPENAI_KEY") or os.environ.get("OPENAI_API_KEY")
-        if not self.api_key:
-            raise ValueError("OPENAI_KEY environment variable is required")
-        
-        # Initialize the graph
-        self.graph = self._build_graph()
-        print("LangGraphAgent initialized successfully.")
-    
-    def _build_graph(self):
-        """Build the LangGraph workflow"""
-        # Initialize OpenAI LLM
+# --------------------------------------------------------------------------- #
+#  Costruzione del graph LangGraph                                             #
+# --------------------------------------------------------------------------- #
+def build_graph(provider: str = "openai"):
+    """Restituisce un graph LangGraph pronto all'uso.
+
+    provider: "openai" (default), "google", "groq", "huggingface"
+    """
+    # --- Selezione LLM ------------------------------------------------------ #
+    if provider == "openai":
+        openai_key = os.getenv("OPENAI_KEY")
+        if not openai_key:
+            raise ValueError(
+                "❌ Environment variable OPENAI_KEY mancante. "
+                "Aggiungi la secret dal tab 'Secrets' dello Space."
+            )
         llm = ChatOpenAI(
-            model="gpt-4-turbo",  # Changed from gpt-4-turbo-preview
+            model_name="gpt-4.1",
             temperature=0,
-            api_key=self.api_key
+            openai_api_key=openai_key,
+        )
+
+    elif 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,
+            )
         )
-        
-        # Bind tools to LLM
-        llm_with_tools = llm.bind_tools(tools)
-        
-        # System message
-        sys_msg = SystemMessage(content=system_prompt)
-        
-        # Node functions
-        def assistant(state: MessagesState):
-            """Assistant node"""
-            # Ensure system message is included
-            messages = state["messages"]
-            if not any(isinstance(msg, SystemMessage) for msg in messages):
-                messages = [sys_msg] + messages
-            
-            response = llm_with_tools.invoke(messages)
-            return {"messages": [response]}
-        
-        # Build the graph
-        builder = StateGraph(MessagesState)
-        
-        # Add nodes
-        builder.add_node("assistant", assistant)
-        builder.add_node("tools", ToolNode(tools))
-        
-        # Add edges
-        builder.add_edge(START, "assistant")
-        builder.add_conditional_edges(
-            "assistant",
-            tools_condition,
+    else:
+        raise ValueError(
+            "Invalid provider. Choose 'openai', 'google', 'groq' or 'huggingface'."
         )
-        builder.add_edge("tools", "assistant")
-        
-        # Compile and return
-        return builder.compile()
-    
-    def __call__(self, question: str) -> str:
-        """
-        Process a question and return an answer.
-        
-        Args:
-            question: The question to answer
-            
-        Returns:
-            str: The answer to the question
-        """
-        print(f"Agent received question (first 100 chars): {question[:100]}...")
-        
-        try:
-            # Create message
-            messages = [HumanMessage(content=question)]
-            
-            # Invoke the graph
-            result = self.graph.invoke({"messages": messages})
-            
-            # Extract the final answer
-            ai_messages = [msg for msg in result["messages"] if isinstance(msg, AIMessage)]
-            
-            if ai_messages:
-                answer = ai_messages[-1].content
-                print(f"Agent returning answer (first 100 chars): {answer[:100]}...")
-                return answer
-            else:
-                return "I couldn't generate a response. Please try again."
-                
-        except Exception as e:
-            print(f"Error processing question: {e}")
-            return f"Error: {str(e)}"
-
-# For backwards compatibility and testing
-BasicAgent = LangGraphAgent
 
+    # Abilita tool calling
+    llm_with_tools = llm.bind_tools(tools)
+
+    # ------------------------- NODES --------------------------------------- #
+    def assistant(state: MessagesState):
+        """Invoca il modello."""
+        return {"messages": [llm_with_tools.invoke(state["messages"])]}
+
+    def retriever(state: MessagesState):
+        """Aggiunge alla history un Q/A simile come esempio."""
+        similar = vector_store.similarity_search(state["messages"][0].content)
+        if similar:
+            example_msg = HumanMessage(
+                content=(
+                    "Here I provide a similar question and answer for reference:\n\n"
+                    f"{similar[0].page_content}"
+                )
+            )
+            return {"messages": [sys_msg] + state["messages"] + [example_msg]}
+        return {"messages": [sys_msg] + state["messages"]}
+
+    # --------------------------- GRAPH ------------------------------------- #
+    builder = StateGraph(MessagesState)
+    builder.add_node("retriever", retriever)
+    builder.add_node("assistant", assistant)
+    builder.add_node("tools", ToolNode(tools))
+
+    builder.add_edge(START, "retriever")
+    builder.add_edge("retriever", "assistant")
+    builder.add_conditional_edges("assistant", tools_condition)
+    builder.add_edge("tools", "assistant")
+
+    return builder.compile()
+
+
+# --------------------------------------------------------------------------- #
+#  Quick test (python agent.py)                                               #
+# --------------------------------------------------------------------------- #
 if __name__ == "__main__":
-    # Test the agent
-    print("Testing LangGraphAgent...")
-    if not os.environ.get("OPENAI_KEY"):
-        print("Error: OPENAI_KEY environment variable not set")
-        print("Please set it with: export OPENAI_KEY=your-openai-api-key")
-        exit(1)
-        
-    try:
-        agent = LangGraphAgent()
-        test_questions = [
-            "What is 15 * 23?",
-            "Search Wikipedia for information about quantum computing",
-            "What are the latest developments in AI according to recent papers on Arxiv?",
-        ]
-        
-        for question in test_questions:
-            print(f"\nQuestion: {question}")
-            answer = agent(question)
-            print(f"Answer: {answer}")
-            
-    except Exception as e:
-        print(f"Error during testing: {e}")
+    graph = build_graph(provider="openai")
+    question = "When was a picture of St. Thomas Aquinas first added to the Wikipedia page on the Principle of double effect?"
+    msgs = [HumanMessage(content=question)]
+    result = graph.invoke({"messages": msgs})
+    for m in result["messages"]:
+        m.pretty_print()