from claude_agent_sdk import ClaudeAgentOptions, query

messages = []
async for msg in query(
    prompt="Research the latest trends in AI agents and give me a brief summary and relevant citations.",
    options=ClaudeAgentOptions(model=MODEL, allowed_tools=["WebSearch"]),
):
    print_activity(msg)
    messages.append(msg)

def chain(input: str, prompts: list[str]) -> str:
    """Chain multiple LLM calls sequentially, passing results between steps."""
    result = input
    for i, prompt in enumerate(prompts, 1):
        print(f"\nStep {i}:")
        result = llm_call(f"{prompt}\nInput: {result}")
        print(result)
    return result

def parallel(prompt: str, inputs: list[str], n_workers: int = 3) -> list[str]:
    """Process multiple inputs concurrently with the same prompt."""
    with ThreadPoolExecutor(max_workers=n_workers) as executor:
        futures = [executor.submit(llm_call, f"{prompt}\nInput: {x}") for x in inputs]
        return [f.result() for f in futures]

def route(input: str, routes: dict[str, str]) -> str:
    selector_prompt = f"""
    Analyze the input and select the most appropriate support team
    from these options: {list(routes.keys())}
    First explain your reasoning, then provide your selection in this XML format:
    <reasoning>...</reasoning>
    <selection>The chosen team name</selection>
    Input: {input}""".strip()

    route_response = llm_call(selector_prompt)
    route_key = extract_xml(route_response, "selection").strip().lower()

    selected_prompt = routes[route_key]
    return llm_call(f"{selected_prompt}\nInput: {input}")

ORCHESTRATOR_PROMPT = """
Analyze this task and break it down into 2-3 distinct approaches:

Task: {task}

Return your response in this format:

<analysis>
Explain your understanding of the task and which variations would be valuable.
</analysis>

<tasks>
    <task>
    <type>formal</type>
    <description>Write a precise, technical version that emphasizes specifications</description>
    </task>
    <task>
    <type>conversational</type>
    <description>Write an engaging, friendly version that connects with readers</description>
    </task>
</tasks>
"""

class FlexibleOrchestrator:
    def process(self, task: str, context: dict | None = None) -> dict:
        # Step 1: Get orchestrator response
        orchestrator_input = self._format_prompt(self.orchestrator_prompt, task=task, **context)
        orchestrator_response = llm_call(orchestrator_input, model=self.model)

        # Parse orchestrator response
        analysis = extract_xml(orchestrator_response, "analysis")
        tasks_xml = extract_xml(orchestrator_response, "tasks")
        tasks = parse_tasks(tasks_xml)

        # Step 2: Process each task
        worker_results = []
        for i, task_info in enumerate(tasks, 1):
            worker_input = self._format_prompt(self.worker_prompt, original_task=task,
                task_type=task_info["type"], task_description=task_info["description"])
            worker_response = llm_call(worker_input, model=self.model)
            worker_results.append({"type": task_info["type"], "result": worker_response})

        return {"analysis": analysis, "worker_results": worker_results}

def loop(task, evaluator_prompt, generator_prompt):
    memory = []
    chain_of_thought = []

    thoughts, result = generate(generator_prompt, task)
    memory.append(result)
    chain_of_thought.append({"thoughts": thoughts, "result": result})

    while True:
        evaluation, feedback = evaluate(evaluator_prompt, result, task)
        if evaluation == "PASS":
            return result, chain_of_thought

        context = "\n".join(["Previous attempts:",
            *[f"- {m}" for m in memory], f"\nFeedback: {feedback}"])

        thoughts, result = generate(generator_prompt, task, context)
        memory.append(result)

tools = [
    {
        "name": "calculator",
        "description": "A simple calculator that performs basic arithmetic operations.",
        "input_schema": {
            "type": "object",
            "properties": {
                "expression": {
                    "type": "string",
                    "description": "The mathematical expression to evaluate (e.g., '2 + 3 * 4').",
                }
            },
            "required": ["expression"],
        },
    }
]

message = client.messages.create(
    model=MODEL_NAME,
    messages=[{"role": "user", "content": user_message}],
    tools=tools,
)

if message.stop_reason == "tool_use":
    tool_use = next(b for b in message.content if b.type == "tool_use")
    tool_result = process_tool_call(tool_use.name, tool_use.input)

    response = client.messages.create(
        model=MODEL_NAME,
        messages=[
            {"role": "user", "content": user_message},
            {"role": "assistant", "content": message.content},
            {"role": "user", "content": [{
                "type": "tool_result",
                "tool_use_id": tool_use.id,
                "content": tool_result,
            }],
        }],
        tools=tools,
    )

from claude_agent_sdk import ClaudeAgentOptions, query

messages = []
async for msg in query(
    prompt="Research the latest trends in AI agents and give me a brief summary and relevant citations.",
    options=ClaudeAgentOptions(model=MODEL, allowed_tools=["WebSearch"]),
):
    print_activity(msg)
    messages.append(msg)

async with ClaudeSDKClient(
    options=ClaudeAgentOptions(
        model=MODEL,
        cwd="chief_of_staff_agent",  # Points to subdirectory with our CLAUDE.md
        setting_sources=["project"],
    )
) as agent:
    await agent.query("What's our current runway?")
    async for msg in agent.receive_response():
        print_activity(msg)
        messages.append(msg)

response = client.messages.create(
    model="claude-sonnet-4-6",
    max_tokens=4000,
    thinking={
        "type": "enabled",
        "budget_tokens": 2000
    },
    messages=[{
        "role": "user",
        "content": "Solve this puzzle: Three people check into a hotel..."
    }]
)

write_response = client.messages.create(
    model=MODEL_NAME,
    max_tokens=300,
    cache_control={"type": "ephemeral"},  # <-- one-line change
    messages=[
        {
            "role": "user",
            "content": str(TIMESTAMP)
            + "<book>"
            + book_content   # 整本《傲慢与偏见》约 187k token
            + "</book>"
            + "\n\nWhat is the title of this book? Only output the title.",
        }
    ],
)