Update crewai_agent.py

crewai_agent.py

@@ -1,187 +1,182 @@
 import os
-from typing import List, Dict, Any
+from typing import List, Dict, Any, Optional
 from dotenv import load_dotenv
 
-from crewai import Agent, Task, Crew, Process
-from crewai.tools import BaseTool
+from langchain.agents import AgentType, initialize_agent, Tool
+from langchain.memory import ConversationBufferMemory
+from langchain.schema import BaseMessage, HumanMessage, AIMessage
 from langchain_google_genai import ChatGoogleGenerativeAI
 from langchain_huggingface import ChatHuggingFace, HuggingFaceEndpoint
 from langchain_community.tools.tavily_search import TavilySearchResults
 from langchain_community.document_loaders import WikipediaLoader, ArxivLoader
-from pydantic import BaseModel, Field
+from langchain.tools import tool
+from langchain.prompts import PromptTemplate
+from langchain.chains import LLMChain
 
 # Load environment variables
 load_dotenv()
 
-class CalculatorTool(BaseTool):
-    """Mathematical calculator tool for basic arithmetic operations."""
-    name: str = "calculator"
-    description: str = "Perform basic mathematical operations: add, subtract, multiply, divide, modulus"
+@tool
+def calculator_tool(operation: str, a: float, b: float) -> str:
+    """Perform basic mathematical operations: add, subtract, multiply, divide, modulus
     
-    def _run(self, operation: str, a: float, b: float) -> float:
-        """Execute mathematical operations."""
-        try:
-            if operation == "add":
-                return a + b
-            elif operation == "subtract":
-                return a - b
-            elif operation == "multiply":
-                return a * b
-            elif operation == "divide":
-                if b == 0:
-                    return "Error: Cannot divide by zero"
-                return a / b
-            elif operation == "modulus":
-                return a % b
-            else:
-                return "Error: Unsupported operation"
-        except Exception as e:
-            return f"Error: {str(e)}"
+    Args:
+        operation: The operation to perform (add, subtract, multiply, divide, modulus)
+        a: First number
+        b: Second number
+    
+    Returns:
+        Result of the mathematical operation
+    """
+    try:
+        if operation == "add":
+            return str(a + b)
+        elif operation == "subtract":
+            return str(a - b)
+        elif operation == "multiply":
+            return str(a * b)
+        elif operation == "divide":
+            if b == 0:
+                return "Error: Cannot divide by zero"
+            return str(a / b)
+        elif operation == "modulus":
+            return str(a % b)
+        else:
+            return "Error: Unsupported operation. Use: add, subtract, multiply, divide, modulus"
+    except Exception as e:
+        return f"Error: {str(e)}"
 
-class WikipediaSearchTool(BaseTool):
-    """Wikipedia search tool for research."""
-    name: str = "wikipedia_search"
-    description: str = "Search Wikipedia for information on any topic"
+@tool
+def wikipedia_search_tool(query: str) -> str:
+    """Search Wikipedia for information on any topic
     
-    def _run(self, query: str) -> str:
-        """Search Wikipedia and return formatted results."""
-        try:
-            search_docs = WikipediaLoader(query=query, load_max_docs=2).load()
-            formatted_results = "\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 formatted_results
-        except Exception as e:
-            return f"Error searching Wikipedia: {str(e)}"
+    Args:
+        query: The search query for Wikipedia
+    
+    Returns:
+        Formatted Wikipedia search results
+    """
+    try:
+        search_docs = WikipediaLoader(query=query, load_max_docs=2).load()
+        formatted_results = "\n\n---\n\n".join([
+            f'Source: {doc.metadata["source"]}\nPage: {doc.metadata.get("page", "")}\n\nContent:\n{doc.page_content[:2000]}...'
+            for doc in search_docs
+        ])
+        return formatted_results
+    except Exception as e:
+        return f"Error searching Wikipedia: {str(e)}"
 
-class WebSearchTool(BaseTool):
-    """Web search tool using Tavily."""
-    name: str = "web_search"
-    description: str = "Search the web for current information using Tavily"
+@tool
+def web_search_tool(query: str) -> str:
+    """Search the web for current information using Tavily
     
-    def _run(self, query: str) -> str:
-        """Search the web and return formatted results."""
-        try:
-            search_results = TavilySearchResults(max_results=3).invoke(query)
-            formatted_results = "\n\n---\n\n".join([
-                f'<Document source="{result.get("url", "")}">\n{result.get("content", "")}\n</Document>'
-                for result in search_results
-            ])
-            return formatted_results
-        except Exception as e:
-            return f"Error searching web: {str(e)}"
+    Args:
+        query: The search query for web search
+    
+    Returns:
+        Formatted web search results
+    """
+    try:
+        search_results = TavilySearchResults(max_results=3).invoke(query)
+        formatted_results = "\n\n---\n\n".join([
+            f'Source: {result.get("url", "")}\n\nContent:\n{result.get("content", "")}'
+            for result in search_results
+        ])
+        return formatted_results
+    except Exception as e:
+        return f"Error searching web: {str(e)}"
 
-class ArxivSearchTool(BaseTool):
-    """ArXiv search tool for academic papers."""
-    name: str = "arxiv_search"
-    description: str = "Search ArXiv for academic papers and research"
+@tool
+def arxiv_search_tool(query: str) -> str:
+    """Search ArXiv for academic papers and research
     
-    def _run(self, query: str) -> str:
-        """Search ArXiv and return formatted results."""
-        try:
-            search_docs = ArxivLoader(query=query, load_max_docs=3).load()
-            formatted_results = "\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 formatted_results
-        except Exception as e:
-            return f"Error searching ArXiv: {str(e)}"
+    Args:
+        query: The search query for ArXiv
+    
+    Returns:
+        Formatted ArXiv search results
+    """
+    try:
+        search_docs = ArxivLoader(query=query, load_max_docs=3).load()
+        formatted_results = "\n\n---\n\n".join([
+            f'Source: {doc.metadata["source"]}\nTitle: {doc.metadata.get("Title", "")}\n\nContent:\n{doc.page_content[:1500]}...'
+            for doc in search_docs
+        ])
+        return formatted_results
+    except Exception as e:
+        return f"Error searching ArXiv: {str(e)}"
 
-class CrewAIAgent:
-    """Multi-purpose CrewAI agent with various capabilities."""
+class LangChainAgent:
+    """Multi-purpose LangChain agent with various capabilities."""
     
     def __init__(self, provider: str = "google"):
-        """Initialize the CrewAI agent with specified LLM provider."""
+        """Initialize the LangChain agent with specified LLM provider."""
         self.provider = provider
         self.llm = self._get_llm(provider)
         self.tools = self._initialize_tools()
-        self.agents = self._create_agents()
+        self.memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
+        self.agent = self._create_agent()
         
     def _get_llm(self, provider: str):
         """Get the specified LLM."""
         if provider == "google":
-            return ChatGoogleGenerativeAI(model="gemini-2.0-flash", temperature=0)
+            return ChatGoogleGenerativeAI(
+                model="gemini-2.0-flash", 
+                temperature=0,
+                max_tokens=2048
+            )
         elif provider == "huggingface":
             return ChatHuggingFace(
                 llm=HuggingFaceEndpoint(
                     repo_id="microsoft/DialoGPT-medium",
                     temperature=0,
+                    max_length=2048,
                 ),
             )
         else:
             raise ValueError("Invalid provider. Choose 'google' or 'huggingface'.")
     
-    def _initialize_tools(self) -> List[BaseTool]:
+    def _initialize_tools(self) -> List[Tool]:
         """Initialize all available tools."""
         return [
-            CalculatorTool(),
-            WikipediaSearchTool(),
-            WebSearchTool(),
-            ArxivSearchTool(),
+            calculator_tool,
+            wikipedia_search_tool,
+            web_search_tool,
+            arxiv_search_tool,
         ]
     
-    def _create_agents(self) -> Dict[str, Agent]:
-        """Create specialized agents for different tasks."""
-        
-        # Research Agent
-        research_agent = Agent(
-            role='Research Specialist',
-            goal='Gather comprehensive and accurate information from multiple sources',
-            backstory="""You are an expert researcher with access to Wikipedia, ArXiv, and web search tools. 
-            You excel at finding relevant, current, and reliable information on any topic.""",
-            tools=[tool for tool in self.tools if 'search' in tool.name],
-            llm=self.llm,
-            verbose=True,
-            allow_delegation=False
-        )
-        
-        # Calculation Agent
-        calculation_agent = Agent(
-            role='Mathematical Analyst',
-            goal='Perform accurate mathematical calculations and analysis',
-            backstory="""You are a mathematical expert capable of performing various calculations 
-            and explaining mathematical concepts clearly.""",
-            tools=[tool for tool in self.tools if 'calculator' in tool.name],
-            llm=self.llm,
-            verbose=True,
-            allow_delegation=False
-        )
-        
-        # General Assistant Agent
-        general_agent = Agent(
-            role='General Assistant',
-            goal='Provide comprehensive answers by coordinating with specialized agents',
-            backstory="""You are a versatile AI assistant that can handle various types of questions 
-            by leveraging specialized tools and knowledge.""",
+    def _create_agent(self):
+        """Create the LangChain agent with tools."""
+        return initialize_agent(
             tools=self.tools,
             llm=self.llm,
+            agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
+            memory=self.memory,
             verbose=True,
-            allow_delegation=True
+            handle_parsing_errors=True,
+            max_iterations=5,
+            early_stopping_method="generate"
         )
-        
-        return {
-            'research': research_agent,
-            'calculation': calculation_agent,
-            'general': general_agent
-        }
     
-    def _determine_agent_type(self, question: str) -> str:
-        """Determine which agent is best suited for the question."""
+    def _determine_approach(self, question: str) -> str:
+        """Determine the best approach for answering the question."""
         question_lower = question.lower()
         
         # Check for mathematical operations
-        math_keywords = ['calculate', 'compute', 'add', 'subtract', 'multiply', 'divide', 'math', 'equation']
+        math_keywords = ['calculate', 'compute', 'add', 'subtract', 'multiply', 'divide', 'math', 'equation', '+', '-', '*', '/', '%']
         if any(keyword in question_lower for keyword in math_keywords):
             return 'calculation'
         
         # Check for research-related queries
-        research_keywords = ['search', 'find', 'research', 'information', 'what is', 'who is', 'when', 'where', 'how']
+        research_keywords = ['search', 'find', 'research', 'information', 'what is', 'who is', 'when', 'where', 'how', 'why']
         if any(keyword in question_lower for keyword in research_keywords):
             return 'research'
         
-        # Default to general agent
+        # Check for academic/scientific queries
+        academic_keywords = ['paper', 'study', 'research', 'academic', 'scientific', 'arxiv', 'journal']
+        if any(keyword in question_lower for keyword in academic_keywords):
+            return 'academic'
+        
         return 'general'
     
     def __call__(self, question: str) -> str:
@@ -189,61 +184,71 @@ class CrewAIAgent:
         try:
             print(f"Processing question: {question[:100]}...")
             
-            # Determine the best agent for this question
-            agent_type = self._determine_agent_type(question)
-            selected_agent = self.agents[agent_type]
-            
-            print(f"Selected agent: {agent_type}")
+            # Determine the best approach for this question
+            approach = self._determine_approach(question)
+            print(f"Selected approach: {approach}")
             
-            # Create a task for the selected agent
-            task = Task(
-                description=f"""
-                Answer the following question comprehensively and accurately:
+            # Create a comprehensive prompt based on the approach
+            if approach == 'calculation':
+                enhanced_question = f"""
+                Solve this mathematical problem step by step:
                 
-                Question: {question}
+                {question}
                 
-                Guidelines:
-                - Use appropriate tools when needed
-                - Provide detailed and helpful responses
-                - Cite sources when using external information
-                - Show calculations step by step for mathematical problems
-                - Be clear and concise in your explanations
-                """,
-                agent=selected_agent,
-                expected_output="A comprehensive and accurate answer to the user's question"
-            )
-            
-            # Create and execute the crew
-            crew = Crew(
-                agents=[selected_agent],
-                tasks=[task],
-                process=Process.sequential,
-                verbose=True
-            )
-            
-            # Execute the task
-            result = crew.kickoff()
-            
-            # Extract the answer from the result
-            if hasattr(result, 'raw'):
-                answer = result.raw
-            elif isinstance(result, str):
-                answer = result
+                Use the calculator tool if needed for complex calculations. Show your work clearly.
+                """
+            elif approach == 'research':
+                enhanced_question = f"""
+                Research and provide comprehensive information about:
+                
+                {question}
+                
+                Use Wikipedia search and web search tools to gather current and accurate information.
+                Cite your sources and provide detailed explanations.
+                """
+            elif approach == 'academic':
+                enhanced_question = f"""
+                Find academic and scientific information about:
+                
+                {question}
+                
+                Use ArXiv search and other research tools to find relevant academic papers and studies.
+                Provide citations and summarize key findings.
+                """
             else:
-                answer = str(result)
+                enhanced_question = f"""
+                Provide a comprehensive answer to:
+                
+                {question}
+                
+                Use appropriate tools as needed (calculator, search tools) to provide accurate information.
+                """
             
-            print(f"Generated answer: {answer[:200]}...")
-            return answer
+            # Use the agent to process the question
+            result = self.agent.run(enhanced_question)
+            
+            print(f"Generated answer: {str(result)[:200]}...")
+            return str(result)
             
         except Exception as e:
             error_msg = f"Error processing question: {str(e)}"
             print(error_msg)
-            return error_msg
+            # Provide a fallback response
+            try:
+                # Try a simple LLM response without tools
+                fallback_result = self.llm.invoke([HumanMessage(content=question)])
+                return fallback_result.content
+            except Exception as fallback_error:
+                return f"Error: Unable to process question. {str(e)}"
+
+    def reset_memory(self):
+        """Reset the conversation memory."""
+        self.memory.clear()
 
 # Test function
-def test_crewai_agent():
-    """Test the CrewAI agent with sample questions."""
-    agent = CrewAIAgent(provider="google")
+def test_langchain_agent():
+    """Test the LangChain agent with sample questions."""
+    agent = LangChainAgent(provider="google")
     
     test_questions = [
         "What is 25 * 34?",
@@ -257,6 +262,7 @@ def test_crewai_agent():
         answer = agent(question)
         print(f"Answer: {answer}")
         print("-" * 50)
+        agent.reset_memory()  # Reset memory between questions for testing
 
 if __name__ == "__main__":
-    test_crewai_agent()
+    test_langchain_agent()