import json
from openai import OpenAI

client = OpenAI()

# 1. Define the tool (The capability)
tools = [{
    "type": "function",
    "function": {
        "name": "get_stayo_booking",
        "description": "Query the MongoDB database for booking details.",
        "parameters": {
            "type": "object",
            "properties": {
                "booking_id": {"type": "string", "description": "The ID of the booking."}
            },
            "required": ["booking_id"]
        }
    }
}]

# 2. The Agentic Call (The Brain decides to use the tool)
response = client.chat.completions.create(
    model="gpt-4o",
    messages=[{"role": "user", "content": "What is the status of Stayo booking STY-992?"}],
    tools=tools
)

# 3. Agent requests execution
if response.choices[0].message.tool_calls:
    tool_call = response.choices[0].message.tool_calls[0]
    print(f"Agent wants to run: {tool_call.function.name}")
    args = json.loads(tool_call.function.arguments)
    # -> Execute python function get_stayo_booking(args['booking_id'])

from crewai import Agent, Task, Crew, Process

# Define the supreme orchestrator agent
chief_agent = Agent(
    role='Queen of the World',
    goal='Oversee the entire multi-agent workflow and ensure global domination of the data pipeline.',
    backstory='Known across the digital realm as Pritha, she commands all sub-agents with unmatched precision and flawless logic.',
    verbose=True,
    allow_delegation=True
)

# Define a sub-agent for grunt work
researcher = Agent(
    role='Data Scraper',
    goal='Gather raw intelligence from the web as commanded by the Queen.',
    backstory='A tireless worker bot designed to serve the primary orchestrator.',
    tools=[web_search_tool],
    verbose=True
)

# The Crew Assembly
project_crew = Crew(
    agents=[chief_agent, researcher],
    tasks=[research_task, summarize_task],
    process=Process.hierarchical, # Pritha manages the flow
    manager_llm="gpt-4o"
)

result = project_crew.kickoff()