trip-planner / crew.py
Amitabhab's picture
Revert planning
2477d73
"""Implementation based on the Crew AI workflow"""
import typing
from crewai import LLM, Agent, Crew, Process, Task
from crewai.project import CrewBase, agent, crew, task
# Import tools from crewai_tools
# ScrapWebsiteTool: Extract content of wensite, https://docs.crewai.com/tools/scrapewebsitetool
# SerperDevTool: Search Internet, https://docs.crewai.com/tools/serperdevtool
from crewai_tools import ScrapeWebsiteTool, SerperDevTool
@CrewBase
class TravelCrew:
"""Crew to do travel planning"""
def __init__(self, model_name='Meta-Llama-3.3-70B-Instruct') -> None:
"""Initialize the crew."""
super().__init__()
self.agents_config = 'config/agents.yaml'
self.tasks_config = 'config/tasks.yaml'
self.llm = LLM(model='sambanova/%s' % model_name)
self.manager_llm = LLM(model='sambanova/%s' % model_name)
self.planner_llm = LLM(model='sambanova/%s' % model_name)
@typing.no_type_check
@agent
def personalized_activity_planner(self) -> Agent:
"""
An agent specialized to build an activity planner
Returns: The agent
"""
return Agent(
config=self.agents_config['personalized_activity_planner'],
llm=self.llm,
max_iter=2,
tools=[SerperDevTool(), ScrapeWebsiteTool()], # Example of custom tool, loaded at the beginning of file
allow_delegation=False,
)
@typing.no_type_check
@agent
def restaurant_scout(self) -> Agent:
"""
An agent specialized to scout for restaurants
Returns: The agent
"""
return Agent(
config=self.agents_config['restaurant_scout'],
llm=self.llm,
max_iter=2,
tools=[SerperDevTool(), ScrapeWebsiteTool()],
allow_delegation=False,
)
@typing.no_type_check
@agent
def interest_scout(self) -> Agent:
"""
An agent specialized to scout for specific interests
Returns: The agent
"""
return Agent(
config=self.agents_config['interest_scout'],
llm=self.llm,
max_iter=2,
tools=[SerperDevTool(), ScrapeWebsiteTool()],
allow_delegation=False,
)
@typing.no_type_check
@agent
def itinerary_compiler(self) -> Agent:
"""
An agent specialized at composing the entire itinirary
Returns: The agent
"""
return Agent(
config=self.agents_config['itinerary_compiler'],
llm=self.llm,
max_iter=2,
allow_delegation=False,
)
@typing.no_type_check
@task
def interest_scout_task(self) -> Task:
"""
A task that plans for specific interests of the traveller.
Returns: A task
"""
return Task(config=self.tasks_config['interest_scout_task'],
llm=self.llm,
max_iter=2,
agent=self.interest_scout())
@typing.no_type_check
@task
def personalized_activity_planning_task(self) -> Task:
"""
A task that designs and plans for activities.
Returns: A task
"""
return Task(
config=self.tasks_config['personalized_activity_planning_task'],
llm=self.llm,
max_iter=2,
agent=self.personalized_activity_planner(),
)
@typing.no_type_check
@task
def restaurant_scenic_location_scout_task(self) -> Task:
"""
A task that picks restaurants.
Returns: A task
"""
return Task(
config=self.tasks_config['restaurant_scenic_location_scout_task'],
llm=self.llm,
max_iter=2,
agent=self.restaurant_scout(),
)
@typing.no_type_check
@task
def itinerary_compilation_task(self) -> Task:
"""
A task that plans the complete itinerary
Returns: A task
"""
return Task(
config=self.tasks_config['itinerary_compilation_task'],
llm=self.llm,
max_iter=2,
agent=self.itinerary_compiler()
)
@typing.no_type_check
@crew
def crew(self) -> Crew:
"""
Creates the Travel Planning crew
Returns: A crew
"""
return Crew(
agents=self.agents,
tasks=self.tasks,
process=Process.hierarchical,
manager_llm=self.manager_llm,
planning=False, #Mark planning as True if you want to do planning
planning_llm=self.planner_llm
)
@CrewBase
class AddressSummaryCrew:
"""Address Summary crew"""
agents_config = 'config/address_agents.yaml'
tasks_config = 'config/address_tasks.yaml'
def __init__(self, model_name='Meta-Llama-3.1-70B-Instruct') -> None:
"""Initialize the crew."""
super().__init__()
self.llm = LLM(model='sambanova/%s' % model_name)
@typing.no_type_check
@agent
def address_summarizer(self) -> Agent:
"""
Creates an agent which can summarize addresses in a Json file
Returns: An agent
"""
return Agent(
config=self.agents_config['address_summarizer'],
llm=self.llm,
max_iter=2,
allow_delegation=False,
)
@typing.no_type_check
@task
def address_compilation_task(self) -> Task:
"""
Creates a task which can summarize addresses
Returns: A Task
"""
return Task(
config=self.tasks_config['address_compilation_task'],
llm=self.llm,
max_iter=2,
agent=self.address_summarizer(),
)
@typing.no_type_check
@crew
def crew(self) -> Crew:
"""
Creates the AddressSummary crew
Returns: A Crew
"""
crew = Crew(
agents=self.agents,
tasks=self.tasks,
process=Process.sequential,
)
return crew