Update crewai_agent.py

crewai_agent.py

@@ -6,10 +6,11 @@ from dotenv import load_dotenv
 load_dotenv()
 
 from langchain.agents import AgentType, initialize_agent, Tool
-from langchain.memory import ConversationBufferMemory
+from langchain.memory import ConversationBufferWindowMemory, ConversationSummaryBufferMemory
 from langchain_core.messages import BaseMessage, HumanMessage, AIMessage
 from langchain_google_genai import ChatGoogleGenerativeAI
 from langchain_huggingface import ChatHuggingFace, HuggingFaceEndpoint
+from langchain_groq import ChatGroq
 from langchain_community.tools.tavily_search import TavilySearchResults
 from langchain_community.document_loaders import WikipediaLoader, ArxivLoader
 from langchain_core.tools import tool
@@ -19,36 +20,59 @@ from langchain.chains import LLMChain
 # Load environment variables
 GOOGLE_API_KEY = os.getenv('GOOGLE_API_KEY')
 HUGGINGFACE_API_TOKEN = os.getenv('HUGGINGFACE_API_TOKEN')
+GROQ_API_KEY = os.getenv('GROQ_API_KEY')
 TAVILY_API_KEY = os.getenv('TAVILY_API_KEY')
 
 
 @tool
-def calculator_tool(operation: str, a: float, b: float) -> str:
-    """Perform basic mathematical operations: add, subtract, multiply, divide, modulus
+def calculator_tool(expression: str) -> str:
+    """Perform mathematical calculations and evaluate expressions
     
     Args:
-        operation: The operation to perform (add, subtract, multiply, divide, modulus)
-        a: First number
-        b: Second number
+        expression: A mathematical expression to evaluate (e.g., "2+2", "25*34", "sqrt(16)", "sin(0.5)")
     
     Returns:
-        Result of the mathematical operation
+        Result of the mathematical expression
     """
     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)
+        import math
+        import re
+        
+        # Clean the expression and make it safe
+        expression = expression.strip()
+        
+        # Replace common mathematical functions
+        expression = expression.replace('sqrt', 'math.sqrt')
+        expression = expression.replace('sin', 'math.sin')
+        expression = expression.replace('cos', 'math.cos')
+        expression = expression.replace('tan', 'math.tan')
+        expression = expression.replace('log', 'math.log')
+        expression = expression.replace('ln', 'math.log')
+        expression = expression.replace('log10', 'math.log10')
+        expression = expression.replace('pi', 'math.pi')
+        expression = expression.replace('e', 'math.e')
+        expression = expression.replace('^', '**')  # Python uses ** for power
+        expression = expression.replace('pow', '**')
+        
+        # More comprehensive regex for mathematical expressions
+        safe_pattern = r'^[0-9+\-*/.() mathsqrtsincolgtanpienpow]+$'
+        if re.match(safe_pattern, expression.replace(' ', '')):
+            # Create a safe namespace for eval
+            safe_dict = {
+                "__builtins__": {},
+                "math": math,
+                "abs": abs,
+                "round": round,
+                "min": min,
+                "max": max
+            }
+            result = eval(expression, safe_dict)
+            return str(result)
         else:
-            return "Error: Unsupported operation. Use: add, subtract, multiply, divide, modulus"
+            return "Error: Invalid characters in expression. Use only numbers and basic math operations."
+            
+    except ZeroDivisionError:
+        return "Error: Cannot divide by zero"
     except Exception as e:
         return f"Error: {str(e)}"
 
@@ -118,17 +142,34 @@ def arxiv_search_tool(query: str) -> str:
 class LangChainAgent:
     """Multi-purpose LangChain agent with various capabilities."""
     
-    def __init__(self, provider: str = "google"):
+    def __init__(self, provider: str = "groq"):
         """Initialize the LangChain agent with specified LLM provider."""
         self.provider = provider
         self.llm = self._get_llm(provider)
         self.tools = self._initialize_tools()
-        self.memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
+        # Use ConversationSummaryBufferMemory for better long-term memory management
+        self.memory = ConversationSummaryBufferMemory(
+            llm=self.llm,
+            memory_key="chat_history",
+            return_messages=True,
+            max_token_limit=2000,  # Limit memory to prevent token overflow
+            moving_summary_buffer="The human and AI are having a conversation about various topics."
+        )
         self.agent = self._create_agent()
         
     def _get_llm(self, provider: str):
         """Get the specified LLM."""
-        if provider == "google":
+        if provider == "groq":
+            if not GROQ_API_KEY:
+                raise ValueError("GROQ_API_KEY not found in environment variables")
+            return ChatGroq(
+                model="llama-3.3-70b-versatile",  # Latest Llama model available on Groq
+                temperature=0.1,
+                max_tokens=8192,  # Increased token limit for better responses
+                api_key=GROQ_API_KEY,
+                streaming=False
+            )
+        elif provider == "google":
             if not GOOGLE_API_KEY:
                 raise ValueError("GOOGLE_API_KEY not found in environment variables")
             return ChatGoogleGenerativeAI(
@@ -149,7 +190,7 @@ class LangChainAgent:
                 ),
             )
         else:
-            raise ValueError("Invalid provider. Choose 'google' or 'huggingface'.")
+            raise ValueError("Invalid provider. Choose 'groq', 'google' or 'huggingface'.")
     
     def _initialize_tools(self) -> List[Tool]:
         """Initialize all available tools."""
@@ -170,8 +211,9 @@ class LangChainAgent:
                 memory=self.memory,
                 verbose=True,
                 handle_parsing_errors=True,
-                max_iterations=3,
-                early_stopping_method="generate"
+                max_iterations=4,
+                early_stopping_method="generate",
+                return_intermediate_steps=False
             )
         except Exception as e:
             print(f"Error creating agent: {e}")
@@ -220,37 +262,56 @@ class LangChainAgent:
             # Create a comprehensive prompt based on the approach
             if approach == 'calculation':
                 enhanced_question = f"""
-                Solve this mathematical problem step by step:
-                
-                {question}
-                
-                Use the calculator tool if needed for complex calculations. Show your work clearly.
+You are a mathematical assistant. Solve this problem step by step:
+
+{question}
+
+IMPORTANT: Use the calculator_tool for ALL mathematical calculations, even simple ones.
+Examples:
+- For "25 * 34", use: calculator_tool("25 * 34")
+- For "sqrt(16)", use: calculator_tool("sqrt(16)")
+- For "2 + 2", use: calculator_tool("2 + 2")
+
+Always show your work and use the tools provided.
                 """
             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.
+You are a research assistant. Provide comprehensive information about:
+
+{question}
+
+IMPORTANT: Use the appropriate search tools to gather information:
+- wikipedia_search_tool("your search query") for general knowledge
+- web_search_tool("your search query") for current information
+- arxiv_search_tool("your search query") for academic papers
+
+Always 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.
+You are an academic research assistant. Find scholarly information about:
+
+{question}
+
+IMPORTANT: Use research tools to find information:
+- arxiv_search_tool("your search query") for academic papers
+- wikipedia_search_tool("your search query") for background information
+
+Provide citations and summarize key findings.
                 """
             else:
                 enhanced_question = f"""
-                Provide a comprehensive answer to:
-                
-                {question}
-                
-                Use appropriate tools as needed (calculator, search tools) to provide accurate information.
+You are a helpful assistant. Answer this question comprehensively:
+
+{question}
+
+IMPORTANT: Use the appropriate tools as needed:
+- calculator_tool("expression") for mathematical calculations
+- wikipedia_search_tool("query") for general information
+- web_search_tool("query") for current information
+- arxiv_search_tool("query") for academic research
+
+Always use tools when they can help provide better answers.
                 """
             
             # Use the agent to process the question
@@ -277,22 +338,23 @@ class LangChainAgent:
 # Test function
 def test_langchain_agent():
     """Test the LangChain agent with sample questions."""
-    agent = LangChainAgent(provider="google")
+    print("Testing LangChain Agent with Groq Llama...")
+    agent = LangChainAgent(provider="groq")
     
     test_questions = [
-        "What is 25 * 34?",
-        "Who was Albert Einstein?",
+        "What is 25 * 34 + 100?",
+        "Who was Albert Einstein and what were his major contributions?",
         "Search for recent developments in artificial intelligence",
-        "What is the theory of relativity?"
+        "What is quantum computing?"
     ]
     
     for question in test_questions:
         print(f"\nQuestion: {question}")
+        print("-" * 50)
         answer = agent(question)
         print(f"Answer: {answer}")
-        print("-" * 50)
+        print("=" * 80)
         agent.reset_memory()  # Reset memory between questions for testing
 
 if __name__ == "__main__":
     test_langchain_agent()
-