Create agent.py

agent.py

@@ -0,0 +1,236 @@
+"""LangGraph Agent with OpenAI"""
+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 langchain_core.tools import tool
+
+load_dotenv()
+
+# Tools definition
+@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) -> 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
+
+@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()
+        formatted_search_docs = "\n\n---\n\n".join(
+            [
+                f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content[:2000]}\n</Document>'
+                for doc in search_docs
+            ])
+        return formatted_search_docs
+    except Exception as e:
+        return f"Error searching Wikipedia: {str(e)}"
+
+
+
+@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()
+        formatted_search_docs = "\n\n---\n\n".join(
+            [
+                f'<Document title="{doc.metadata.get("Title", "")}" authors="{doc.metadata.get("Authors", "")}"/>\n{doc.page_content[:1500]}\n</Document>'
+                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 helpful AI assistant with access to various tools for calculations and information retrieval.
+You can perform mathematical operations, search Wikipedia, and search academic papers on Arxiv.
+Always try to provide accurate, concise, and helpful responses based on the tools available to you.
+When searching for information, be thorough but concise in your final answer.
+If a question requires multiple steps or tools, break it down and use the appropriate tools in sequence."""
+
+# Tools list
+tools = [
+    multiply,
+    add,
+    subtract,
+    divide,
+    modulus,
+    wiki_search,
+    arxiv_search,
+]
+
+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
+        llm = ChatOpenAI(
+            model="gpt-4.1",  
+            temperature=0.0,
+            api_key=self.api_key
+        )
+        
+        # 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,
+        )
+        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
+
+if __name__ == "__main__":
+    # Test the agent
+    print("Testing LangGraphAgent...")
+    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}")