Rename main.py to crewai_agent.py

crewai_agent.py

@@ -0,0 +1,262 @@
+import os
+from typing import List, Dict, Any
+from dotenv import load_dotenv
+
+from crewai import Agent, Task, Crew, Process
+from crewai.tools import BaseTool
+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
+
+# 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"
+    
+    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)}"
+
+class WikipediaSearchTool(BaseTool):
+    """Wikipedia search tool for research."""
+    name: str = "wikipedia_search"
+    description: 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)}"
+
+class WebSearchTool(BaseTool):
+    """Web search tool using Tavily."""
+    name: str = "web_search"
+    description: 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)}"
+
+class ArxivSearchTool(BaseTool):
+    """ArXiv search tool for academic papers."""
+    name: str = "arxiv_search"
+    description: 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)}"
+
+class CrewAIAgent:
+    """Multi-purpose CrewAI agent with various capabilities."""
+    
+    def __init__(self, provider: str = "google"):
+        """Initialize the CrewAI agent with specified LLM provider."""
+        self.provider = provider
+        self.llm = self._get_llm(provider)
+        self.tools = self._initialize_tools()
+        self.agents = self._create_agents()
+        
+    def _get_llm(self, provider: str):
+        """Get the specified LLM."""
+        if provider == "google":
+            return ChatGoogleGenerativeAI(model="gemini-2.0-flash", temperature=0)
+        elif provider == "huggingface":
+            return ChatHuggingFace(
+                llm=HuggingFaceEndpoint(
+                    repo_id="microsoft/DialoGPT-medium",
+                    temperature=0,
+                ),
+            )
+        else:
+            raise ValueError("Invalid provider. Choose 'google' or 'huggingface'.")
+    
+    def _initialize_tools(self) -> List[BaseTool]:
+        """Initialize all available tools."""
+        return [
+            CalculatorTool(),
+            WikipediaSearchTool(),
+            WebSearchTool(),
+            ArxivSearchTool(),
+        ]
+    
+    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.""",
+            tools=self.tools,
+            llm=self.llm,
+            verbose=True,
+            allow_delegation=True
+        )
+        
+        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."""
+        question_lower = question.lower()
+        
+        # Check for mathematical operations
+        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']
+        if any(keyword in question_lower for keyword in research_keywords):
+            return 'research'
+        
+        # Default to general agent
+        return 'general'
+    
+    def __call__(self, question: str) -> str:
+        """Process a question and return an answer."""
+        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}")
+            
+            # Create a task for the selected agent
+            task = Task(
+                description=f"""
+                Answer the following question comprehensively and accurately:
+                
+                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
+            else:
+                answer = str(result)
+            
+            print(f"Generated answer: {answer[:200]}...")
+            return answer
+            
+        except Exception as e:
+            error_msg = f"Error processing question: {str(e)}"
+            print(error_msg)
+            return error_msg
+
+# Test function
+def test_crewai_agent():
+    """Test the CrewAI agent with sample questions."""
+    agent = CrewAIAgent(provider="google")
+    
+    test_questions = [
+        "What is 25 * 34?",
+        "Who was Albert Einstein?",
+        "Search for recent developments in artificial intelligence",
+        "What is the theory of relativity?"
+    ]
+    
+    for question in test_questions:
+        print(f"\nQuestion: {question}")
+        answer = agent(question)
+        print(f"Answer: {answer}")
+        print("-" * 50)
+
+if __name__ == "__main__":
+    test_crewai_agent()