CrewAI agent

1"""This module defines the main entry point for the Apify Actor.
2
3Feel free to modify this file to suit your specific needs.
4
5To build Apify Actors, utilize the Apify SDK toolkit, read more at the official documentation:
6https://docs.apify.com/sdk/python
7"""
8
9from __future__ import annotations
10
11import logging
12
13from apify import Actor
14from crewai import Agent, Crew, Task
15
16from src.models import AgentStructuredOutput
17from src.ppe_utils import charge_for_actor_start, charge_for_model_tokens
18from src.tools import tool_calculator_sum, tool_scrape_instagram_profile_posts
19
20
21async def main() -> None:
22    """Main entry point for the Apify Actor.
23
24    This coroutine is executed using `asyncio.run()`, so it must remain an asynchronous function for proper execution.
25    Asynchronous execution is required for communication with the Apify platform, and it also enhances performance in
26    the field of web scraping significantly.
27
28    Raises:
29        ValueError: If the input is missing required attributes.
30    """
31    async with Actor:
32        # Handle input
33        actor_input = await Actor.get_input()
34
35        query = actor_input.get('query')
36        model_name = actor_input.get('modelName', 'gpt-4o-mini')
37        if debug := actor_input.get('debug', False):
38            Actor.log.setLevel(logging.DEBUG)
39        if not query:
40            msg = 'Missing "query" attribute in input!'
41            raise ValueError(msg)
42
43        await charge_for_actor_start()
44
45        # Create a toolkit for the agent
46        tools = [tool_calculator_sum, tool_scrape_instagram_profile_posts]
47
48        # Create an agent
49        # For more information, see https://docs.crewai.com/concepts/agents
50        agent = Agent(
51            role='Helpful agent',
52            goal='Assist users with various tasks.',
53            backstory='I am a helpful agent that can assist you with various tasks.',
54            tools=tools,
55            verbose=debug,
56        )
57
58        # Create a task assigned to the agent
59        # For more information, see https://docs.crewai.com/concepts/tasks
60        task = Task(
61            description=query,
62            expected_output='A helpful response to the user query.',
63            agent=agent,
64            output_pydantic=AgentStructuredOutput,
65        )
66
67        # Create a one-man crew
68        # For more information, see https://docs.crewai.com/concepts/crews
69        crew = Crew(agents=[agent], tasks=[task])
70
71        # Kick off the crew and get the response
72        crew_output = crew.kickoff()
73        raw_response = crew_output.raw
74        response = crew_output.pydantic
75
76        # Charge the user for the tokens used by the model
77        total_tokens = crew_output.token_usage.total_tokens
78        await charge_for_model_tokens(model_name, total_tokens)
79
80        if not response or not raw_response:
81            Actor.log.error('Failed to get a response from the agent!')
82            await Actor.fail(status_message='Failed to get a response from the agent!')
83
84        # Push results to the dataset
85        await Actor.push_data(
86            {
87                'query': query,
88                'response': raw_response,
89                'structured_response': response.dict() if response else {},
90            }
91        )
92        Actor.log.info('Pushed the data into the dataset!')