from typing import List, Dict, Any
from dataclasses import dataclass

@dataclass
class AgentState:
    goal: str
    context: List[Dict[str, Any]]
    memory: List[str]
    current_step: int
    max_steps: int = 10

class BaseAgent:
    def __init__(self, model: str, tools: List[callable]):
        self.model = model
        self.tools = {t.__name__: t for t in tools}
        self.state = None
    
    def plan(self, goal: str) -> List[str]:
        """Break down goal into executable steps"""
        raise NotImplementedError
    
    def execute_step(self, step: str) -> str:
        """Execute a single step using available tools"""
        raise NotImplementedError
    
    def run(self, goal: str) -> str:
        """Main agent loop"""
        self.state = AgentState(goal=goal, context=[], memory=[], current_step=0)
        steps = self.plan(goal)
        
        for step in steps:
            if self.state.current_step >= self.state.max_steps:
                return "Max steps reached"
            result = self.execute_step(step)
            self.state.memory.append(result)
            self.state.current_step += 1
        
        return self.synthesize_result()

class ReActAgent(BaseAgent):
    SYSTEM_PROMPT = """You are a helpful agent. For each query:
    1. Think about what you need to do
    2. Choose a tool to use
    3. Observe the result
    4. Repeat or give final answer
    
    Available tools: {tools}
    
    Format:
    Thought: [your reasoning]
    Action: [tool_name]
    Action Input: [input to tool]
    Observation: [tool output]
    ... (repeat as needed)
    Final Answer: [your response]
    """
    
    def execute_step(self, step: str) -> str:
        response = self.llm.generate(
            self.SYSTEM_PROMPT.format(tools=list(self.tools.keys())),
            step
        )
        
        # Parse and execute tool calls
        if "Action:" in response:
            tool_name, tool_input = self.parse_action(response)
            observation = self.tools[tool_name](tool_input)
            return f"Observation: {observation}"
        
        return response

class PlanExecuteAgent(BaseAgent):
    def plan(self, goal: str) -> List[str]:
        planning_prompt = f"""
        Goal: {goal}
        
        Create a detailed step-by-step plan.
        Each step should be atomic and executable.
        
        Output as JSON:
        steps
        """
        
        response = self.llm.generate(planning_prompt)
        return json.loads(response)["steps"]
    
    def execute_step(self, step: str) -> str:
        # Can use a different, cheaper model for execution
        execution_prompt = f"""
        Execute this step: {step}
        
        Previous results: {self.state.memory[-3:]}
        Available tools: {list(self.tools.keys())}
        """
        return self.executor_llm.generate(execution_prompt)

from crewai import Agent, Task, Crew

# Define specialized agents
researcher = Agent(
    role='Senior Research Analyst',
    goal='Uncover cutting-edge developments in AI',
    backstory='Expert at finding and synthesizing information',
    tools=[search_tool, scrape_tool],
    llm=claude_model
)

writer = Agent(
    role='Tech Writer',
    goal='Create engaging technical content',
    backstory='Skilled at explaining complex topics simply',
    tools=[write_tool],
    llm=gpt_model
)

reviewer = Agent(
    role='Technical Reviewer',
    goal='Ensure accuracy and quality',
    backstory='Perfectionist with deep technical knowledge',
    tools=[fact_check_tool],
    llm=claude_model
)

# Define tasks
research_task = Task(
    description='Research the latest AI agent frameworks',
    agent=researcher,
    expected_output='Comprehensive research report'
)

writing_task = Task(
    description='Write a blog post based on the research',
    agent=writer,
    context=[research_task]
)

review_task = Task(
    description='Review and improve the blog post',
    agent=reviewer,
    context=[writing_task]
)

# Create and run crew
crew = Crew(
    agents=[researcher, writer, reviewer],
    tasks=[research_task, writing_task, review_task],
    verbose=True
)

result = crew.kickoff()

from pydantic import BaseModel, Field
from typing import Literal

class ToolCall(BaseModel):
    tool: Literal["search", "calculate", "write_file"]
    arguments: Dict[str, Any]
    reasoning: str = Field(description="Why this tool was chosen")

class AgentResponse(BaseModel):
    thought: str
    action: Optional[ToolCall] = None
    final_answer: Optional[str] = None

# Force structured output
response = client.chat.completions.create(
    model="gpt-4o",
    messages=messages,
    response_format={
        "type": "json_schema",
        "json_schema": AgentResponse.model_json_schema()
    }
)

class ResilientAgent(BaseAgent):
    def execute_with_retry(self, step: str, max_retries: int = 3) -> str:
        for attempt in range(max_retries):
            try:
                result = self.execute_step(step)
                if self.validate_result(result):
                    return result
                raise ValueError("Invalid result")
            except Exception as e:
                if attempt == max_retries - 1:
                    return self.fallback_strategy(step, e)
                self.state.memory.append(f"Attempt {attempt + 1} failed: {e}")
    
    def fallback_strategy(self, step: str, error: Exception) -> str:
        # Ask human, use simpler approach, or skip
        return self.ask_human_for_help(step, error)

class CostAwareAgent(BaseAgent):
    def __init__(self, *args, budget_usd: float = 1.0, **kwargs):
        super().__init__(*args, **kwargs)
        self.budget = budget_usd
        self.spent = 0.0
    
    def track_cost(self, response) -> float:
        # Approximate costs
        input_cost = response.usage.prompt_tokens * 0.00001
        output_cost = response.usage.completion_tokens * 0.00003
        return input_cost + output_cost
    
    def run(self, goal: str) -> str:
        while self.spent < self.budget:
            # ... agent logic
            self.spent += self.track_cost(response)
        
        return "Budget exhausted"