from crewai import Agent, Task, Crew
from langchain_openai import ChatOpenAI
from a2a.server import A2AServer
 
# Initialize language model
llm = ChatOpenAI(model="gpt-4")
 
# Create the research agent
researcher = Agent(
    role="Travel Researcher",
    goal="Find the most interesting places to visit",
    backstory="You are an expert travel researcher with deep knowledge of destinations worldwide",
    verbose=True,
    llm=llm
)
 
# Create the planning agent
planner = Agent(
    role="Travel Planner",
    goal="Create detailed itineraries based on research",
    backstory="You are an experienced travel planner who creates perfect itineraries",
    verbose=True,
    llm=llm
)
 
# Create tasks for each agent
research_task = Task(
    description="Research {destination} and identify top attractions",
    agent=researcher
)
 
planning_task = Task(
    description="Create a day-by-day itinerary for {destination} based on research",
    agent=planner
)
 
# Create the crew
travel_crew = Crew(
    agents=[researcher, planner],
    tasks=[research_task, planning_task],
    verbose=True
)
 
# Function to handle A2A requests using CrewAI
def process_with_crewai(user_input):
    # Parse destination from user input (simplified)
    if "travel to" in user_input:
        destination = user_input.split("travel to")[1].strip()
    else:
        destination = "Unknown location"
    
    # Execute the crew tasks with the destination
    result = travel_crew.kickoff(inputs={"destination": destination})
    return result
 
# Create A2A server with CrewAI integration
a2a_server = A2AServer(
    agent=process_with_crewai,
    agent_card={
        "name": "CrewAI Travel Assistant",
        "description": "A team of travel agents that help plan your perfect trip",
        "skills": [
            {
                "id": "travel-research",
                "name": "Travel Research",
                "description": "Research travel destinations and attractions"
            },
            {
                "id": "itinerary-planning",
                "name": "Itinerary Planning",
                "description": "Create detailed travel itineraries"
            }
        ]
    }
)
 
# Run the server
if __name__ == "__main__":
    a2a_server.run(host="0.0.0.0", port=8000)