CrewAI Integration

1from crewai import Agent, Task, Crew, LLM
2from crewai_tools import BaseTool
3from cognitora import Cognitora
4from typing import Type
5from pydantic import BaseModel, Field
6
7cognitora = Cognitora(api_key="your-cognitora-api-key")
8
9class CodeExecutionInput(BaseModel):
10    code: str = Field(description="Python code to execute")
11
12class CodeExecutionTool(BaseTool):
13    name: str = "Execute Python Code"
14    description: str = "Execute Python code in a secure sandbox environment"
15    args_schema: Type[BaseModel] = CodeExecutionInput
16
17    def __init__(self):
18        super().__init__()
19        self.session = cognitora.sessions.create(
20            image="docker.io/library/python:3.11-slim",
21            timeout=300,
22            persistent=True
23        )
24
25    def _run(self, code: str) -> str:
26        execution = cognitora.compute.execute(
27            session_id=self.session.id,
28            command=["python", "-c", code]
29        )
30        
31        if execution.exit_code == 0:
32            return f"Success: {execution.stdout}"
33        else:
34            return f"Error: {execution.stderr}"
35
36# Initialize the tool
37execute_python = CodeExecutionTool()
38
39# Define the agent
40python_executor = Agent(
41    role='Python Executor',
42    goal='Execute Python code and return the results',
43    backstory='You are an expert Python programmer capable of executing code and returning results.',
44    tools=[execute_python],
45    llm=LLM(model="gpt-4o")
46)
47
48# Define the task
49execute_task = Task(
50    description="Calculate how many r's are in the word 'strawberry'",
51    agent=python_executor,
52    expected_output="The number of r's in the word 'strawberry'"
53)
54
55# Create the crew
56code_execution_crew = Crew(
57    agents=[python_executor],
58    tasks=[execute_task],
59    verbose=True,
60)
61
62# Run the crew
63result = code_execution_crew.kickoff()
64print(result)