Agent Architecture

"""ReAct Agent — No frameworks, just the pattern."""
import json
import openai

client = openai.OpenAI()

# ─── Define Tools ────────────────────────────────────────────

def search_wikipedia(query: str) -> str:
    """Search Wikipedia and return a summary."""
    import urllib.request, urllib.parse
    url = f"https://en.wikipedia.org/api/rest_v1/page/summary/{urllib.parse.quote(query)}"
    try:
        with urllib.request.urlopen(url) as resp:
            data = json.loads(resp.read())
            return data.get("extract", "No article found.")
    except Exception as e:
        return f"Search failed: {e}"

def calculate(expression: str) -> str:
    """Evaluate a math expression safely."""
    allowed = set("0123456789+-*/.() ")
    if not all(c in allowed for c in expression):
        return "Error: Invalid characters in expression"
    try:
        result = eval(expression)  # Safe due to character whitelist
        return str(result)
    except Exception as e:
        return f"Calculation error: {e}"

# ─── Tool Registry ───────────────────────────────────────────

TOOLS = {
    "search_wikipedia": {
        "fn": search_wikipedia,
        "description": "Search Wikipedia for a topic. Input: search query string.",
    },
    "calculate": {
        "fn": calculate,
        "description": "Evaluate a math expression. Input: expression string like '2 + 2'.",
    },
}

# ─── System Prompt ───────────────────────────────────────────

SYSTEM_PROMPT = """You are a ReAct agent. For each step, output EXACTLY one of:

Thought: [your reasoning about what to do next]
Action: [tool_name]: [input]

OR when you have enough information:

Thought: [final reasoning]
Answer: [your final response to the user]

Available tools:
{tools}

Rules:
- Always start with a Thought
- Only use one Action per step
- Wait for Observation before your next Thought
- When you have enough info, give the final Answer
""".format(tools="\n".join(f"- {name}: {t['description']}" for name, t in TOOLS.items()))

# ─── Agent Loop ──────────────────────────────────────────────

def run_agent(user_query: str, max_iterations: int = 10) -> str:
    """Execute the ReAct loop until answer or max iterations."""
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": user_query},
    ]
    
    for i in range(max_iterations):
        # Get LLM response
        response = client.chat.completions.create(
            model="gpt-4o",
            messages=messages,
            temperature=0,
            max_tokens=500,
        )
        
        llm_output = response.choices[0].message.content.strip()
        messages.append({"role": "assistant", "content": llm_output})
        
        print(f"\n--- Step {i+1} ---")
        print(llm_output)
        
        # Check if agent produced a final answer
        if "Answer:" in llm_output:
            answer = llm_output.split("Answer:")[-1].strip()
            return answer
        
        # Parse and execute action
        if "Action:" in llm_output:
            action_line = llm_output.split("Action:")[-1].strip()
            tool_name, tool_input = action_line.split(":", 1)
            tool_name = tool_name.strip()
            tool_input = tool_input.strip()
            
            if tool_name in TOOLS:
                observation = TOOLS[tool_name]["fn"](tool_input)
            else:
                observation = f"Error: Unknown tool '{tool_name}'"
            
            print(f"Observation: {observation[:200]}")
            messages.append({"role": "user", "content": f"Observation: {observation}"})
        else:
            # No action and no answer — prompt for next step
            messages.append({"role": "user", "content": "Continue with your next Thought or Answer."})
    
    return "Agent reached maximum iterations without a final answer."

# ─── Run It ──────────────────────────────────────────────────

result = run_agent("What is the population of France divided by the area of Germany?")
print(f"\nFinal Answer: {result}")

"""Robust agent with error handling, retries, and budget tracking."""

class AgentBudget:
    """Track and limit agent resource consumption."""
    def __init__(self, max_iterations=10, max_tokens=50000, max_time_seconds=60):
        self.max_iterations = max_iterations
        self.max_tokens = max_tokens
        self.max_time_seconds = max_time_seconds
        self.iterations_used = 0
        self.tokens_used = 0
        self.start_time = None
    
    def check(self):
        import time
        if self.iterations_used >= self.max_iterations:
            raise BudgetExceeded("Max iterations reached")
        if self.tokens_used >= self.max_tokens:
            raise BudgetExceeded("Token budget exhausted")
        if self.start_time and (time.time() - self.start_time) > self.max_time_seconds:
            raise BudgetExceeded("Time limit exceeded")

class BudgetExceeded(Exception):
    pass

def run_robust_agent(query: str, budget: AgentBudget = None) -> dict:
    """Agent with full error handling and observability."""
    import time
    
    budget = budget or AgentBudget()
    budget.start_time = time.time()
    trace = []  # Full execution trace for debugging
    
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": query},
    ]
    
    try:
        while True:
            budget.check()
            budget.iterations_used += 1
            
            response = client.chat.completions.create(
                model="gpt-4o",
                messages=messages,
                temperature=0,
            )
            
            llm_output = response.choices[0].message.content.strip()
            budget.tokens_used += response.usage.total_tokens
            
            trace.append({"step": budget.iterations_used, "output": llm_output})
            messages.append({"role": "assistant", "content": llm_output})
            
            if "Answer:" in llm_output:
                return {"answer": llm_output.split("Answer:")[-1].strip(),
                        "trace": trace, "budget": vars(budget)}
            
            if "Action:" in llm_output:
                action_line = llm_output.split("Action:")[-1].strip()
                try:
                    tool_name, tool_input = action_line.split(":", 1)
                    tool_name = tool_name.strip()
                    observation = TOOLS[tool_name]["fn"](tool_input.strip())
                except KeyError:
                    observation = f"Error: Tool '{tool_name}' not found. Available: {list(TOOLS.keys())}"
                except Exception as e:
                    observation = f"Error executing tool: {e}"
                
                trace.append({"step": budget.iterations_used, "observation": observation})
                messages.append({"role": "user", "content": f"Observation: {observation}"})
    
    except BudgetExceeded as e:
        return {"answer": None, "error": str(e), "trace": trace, "budget": vars(budget)}