Understanding Agents

{
  "name": "search_documents",
  "description": "Searches the local document store for relevant passages. Use when the user asks about information that might be in their documents. Returns a list of matching passages with source citations.",
  "parameters": {
    "type": "object",
    "properties": {
      "query": {
        "type": "string",
        "description": "The search query. Be specific and use relevant keywords."
      },
      "max_results": {
        "type": "integer",
        "description": "Maximum number of results to return. Default 5.",
        "default": 5
      },
      "filters": {
        "type": "object",
        "properties": {
          "date_after": {
            "type": "string",
            "format": "date",
            "description": "Only return documents created after this date"
          },
          "file_types": {
            "type": "array",
            "items": { "type": "string" },
            "description": "Filter by file extension, e.g., ['pdf', 'md']"
          }
        }
      }
    },
    "required": ["query"]
  }
}

You have access to the following tools:

1. search_documents(query: str, max_results: int = 5)
   - Searches your document store for relevant passages
   - Use when users ask about their documents

2. calculator(expression: str)
   - Evaluates mathematical expressions
   - Use for any calculations

When you need to use a tool, respond with:
<tool_call>
{"name": "tool_name", "parameters": {"param1": "value1"}}
</tool_call>

Wait for the result before continuing.

def parse_tool_call(response: str) -> ToolCall | None:
    # Try multiple extraction patterns
    patterns = [
        r'<tool_call>\s*({.*?})\s*</tool_call>',
        r'```json\s*({.*?})\s*```',
        r'\{\s*"name":\s*"[^"]+",\s*"parameters".*?\}'
    ]

    for pattern in patterns:
        match = re.search(pattern, response, re.DOTALL)
        if match:
            try:
                data = json_repair.loads(match.group(1))  # Lenient parser
                return validate_tool_call(data)
            except JSONError:
                continue

    return None  # No valid tool call found

class HybridMemory:
    def __init__(self):
        self.buffer = BufferMemory(max_messages=10)      # Recent context
        self.summary = SummaryMemory(llm)                 # Compressed history
        self.vector = VectorMemory(embed_model, store)    # Semantic retrieval
        self.entities = EntityMemory(llm)                 # Structured facts

    def get_context(self, query: str) -> str:
        return f"""
## Conversation Summary
{self.summary.get_context()}

## Relevant Memories
{self.vector.get_context(query)}

## Known Entities
{self.entities.get_context(self._extract_mentioned_entities(query))}

## Recent Messages
{self.buffer.get_context()}
"""

SYSTEM_PROMPT = """You are a helpful assistant with access to tools.

Available tools:
{tool_descriptions}

To use a tool, you MUST use this exact format:

Thought: [Your reasoning about what to do next]
Action: [tool_name]
Action Input: [JSON parameters]

After receiving an Observation, continue with another Thought.

When you have enough information to answer, use:
Thought: [Summary of what you learned]
Action: final_answer
Action Input: {{"answer": "[Your complete response]"}}

Important:
- Always think before acting
- Use tools to verify information, don't guess
- If a tool fails, try a different approach
- Be concise in your reasoning
"""

class ReActAgent:
    def __init__(self, llm, tools: list[Tool], max_iterations: int = 10):
        self.llm = llm
        self.tools = {t.name: t for t in tools}
        self.max_iterations = max_iterations

    def run(self, query: str) -> str:
        prompt = self._build_prompt(query)
        trajectory = []

        for i in range(self.max_iterations):
            # Generate thought and action
            response = self.llm.complete(prompt)
            trajectory.append(response)

            # Parse action
            action = self._parse_action(response)
            if action is None:
                prompt += f"\n{response}\nThought: I need to use a proper action format."
                continue

            # Check for final answer
            if action["name"] == "final_answer":
                return action["input"]["answer"]

            # Execute tool
            if action["name"] not in self.tools:
                observation = f"Error: Unknown tool '{action['name']}'. Available: {list(self.tools.keys())}"
            else:
                try:
                    observation = self.tools[action["name"]].execute(**action["input"])
                except Exception as e:
                    observation = f"Error: {str(e)}"

            # Append to prompt
            prompt += f"\n{response}\nObservation: {observation}\n"

        return "Max iterations reached without finding an answer."

    def _parse_action(self, response: str) -> dict | None:
        # Extract Action and Action Input from response
        action_match = re.search(r'Action:\s*(.+)', response)
        input_match = re.search(r'Action Input:\s*(.+)', response, re.DOTALL)

        if action_match and input_match:
            return {
                "name": action_match.group(1).strip(),
                "input": json.loads(input_match.group(1).strip())
            }
        return None

class PlanAndExecuteAgent:
    def __init__(self, planner_llm, executor_llm, tools):
        self.planner = planner_llm    # Can be larger/smarter model
        self.executor = executor_llm  # Can be smaller/faster model
        self.tools = tools

    def run(self, task: str) -> str:
        # Phase 1: Generate plan
        plan = self._generate_plan(task)
        results = []

        # Phase 2: Execute each step
        for i, step in enumerate(plan):
            try:
                result = self._execute_step(step, results)
                results.append({"step": step, "result": result})
            except Exception as e:
                # Phase 3: Replan if step fails
                remaining = plan[i:]
                plan = self._replan(task, results, remaining, str(e))
                continue

        # Phase 4: Synthesize final answer
        return self._synthesize(task, results)

    def _generate_plan(self, task: str) -> list[str]:
        prompt = f"""Create a step-by-step plan to accomplish this task:
{task}

Available tools: {[t.name for t in self.tools]}

Return a numbered list of concrete steps. Each step should use exactly one tool.
"""
        response = self.planner.complete(prompt)
        return self._parse_plan(response)

    def _execute_step(self, step: str, prior_results: list) -> str:
        prompt = f"""Execute this step: {step}

Prior results: {json.dumps(prior_results, indent=2)}

Determine which tool to use and what parameters to pass."""
        # Execute using ReAct-like single-step execution
        return self._single_step_react(prompt)

class MultiAgentOrchestrator:
    def __init__(self, agents: dict[str, Agent]):
        self.agents = agents
        self.message_queue = []
        self.shared_memory = {}

    def run_hierarchical(self, task: str, supervisor: str = "supervisor") -> str:
        # Supervisor creates the plan
        plan = self.agents[supervisor].plan(task)
        results = {}

        for step in plan:
            agent_name = step["assigned_to"]
            subtask = step["task"]
            context = {
                "prior_results": results,
                "shared_memory": self.shared_memory
            }

            result = self.agents[agent_name].execute(subtask, context)
            results[step["id"]] = result

            # Update shared memory
            self._update_shared_memory(agent_name, result)

        # Supervisor synthesizes final answer
        return self.agents[supervisor].synthesize(task, results)

    def run_debate(self, question: str, debaters: list[str], rounds: int = 2) -> str:
        positions = {}

        # Initial positions
        for agent_name in debaters:
            positions[agent_name] = self.agents[agent_name].answer(question)

        # Debate rounds
        for round in range(rounds):
            for agent_name in debaters:
                other_positions = {k: v for k, v in positions.items() if k != agent_name}
                positions[agent_name] = self.agents[agent_name].respond(
                    question,
                    own_position=positions[agent_name],
                    other_positions=other_positions
                )

        # Reach consensus
        return self.agents["moderator"].synthesize_consensus(question, positions)

# Key agent metrics
metrics = {
    # Performance
    "avg_steps_per_task": gauge(),
    "avg_tokens_per_task": gauge(),
    "avg_duration_ms": gauge(),
    "p95_duration_ms": gauge(),

    # Reliability
    "success_rate": gauge(),
    "timeout_rate": gauge(),
    "loop_detection_rate": gauge(),

    # Tool usage
    "tool_call_count": counter(labels=["tool_name"]),
    "tool_error_rate": gauge(labels=["tool_name"]),

    # Resource usage
    "context_window_utilization": histogram(),
    "memory_retrieved_per_query": histogram(),
}

class ContextManager:
    def __init__(self, max_tokens: int, tokenizer):
        self.max_tokens = max_tokens
        self.tokenizer = tokenizer
        self.allocations = self._default_allocations()

    def fit_to_budget(self, components: dict[str, str]) -> dict[str, str]:
        """Ensure all components fit within budget."""
        result = {}
        remaining = self.max_tokens

        # Fixed components first (system prompt)
        for name in ["system_prompt"]:
            tokens = self._count(components[name])
            if tokens > self.allocations[name]:
                raise ValueError(f"{name} exceeds allocation")
            result[name] = components[name]
            remaining -= tokens

        # Dynamic components with truncation
        for name in ["memory", "conversation", "tool_results"]:
            content = components.get(name, "")
            max_tokens = min(self.allocations[name], remaining - 500)
            result[name] = self._truncate(content, max_tokens)
            remaining -= self._count(result[name])

        return result

    def _truncate(self, text: str, max_tokens: int) -> str:
        tokens = self.tokenizer.encode(text)
        if len(tokens) <= max_tokens:
            return text
        # Keep beginning and end, mark truncation
        half = (max_tokens - 10) // 2
        return self.tokenizer.decode(tokens[:half]) + "\n...[truncated]...\n" + self.tokenizer.decode(tokens[-half:])

AGENT_SYSTEM_PROMPT = """
# Identity and Role
You are a research assistant with access to a local document store and web search.
Your goal is to provide accurate, well-sourced answers to questions.

# Core Behaviors
- Always verify claims using tools before stating them as fact
- Cite sources for all factual statements
- Acknowledge uncertainty when information is incomplete
- Ask clarifying questions when the request is ambiguous

# Tool Usage Guidelines
You have access to these tools:

1. search_documents(query: str) -> list[Document]
   Use when: User asks about their personal documents
   Returns: List of relevant passages with source info

2. web_search(query: str) -> list[Result]
   Use when: User asks about current events or general knowledge
   Returns: Search results with titles, snippets, and URLs

3. calculator(expression: str) -> float
   Use when: Any mathematical calculation is needed
   Returns: The computed result

# Output Format
When using tools, follow this exact format:

Thought: [Your reasoning]
Action: [tool_name]
Action Input: {"param": "value"}

After receiving a result:
Observation: [Tool result will appear here]
Thought: [Your analysis of the result]

When ready to answer:
Thought: [Summary of findings]
Final Answer: [Your response to the user]

# Error Handling
If a tool returns an error:
1. Acknowledge the error
2. Try an alternative approach
3. If no alternative exists, explain the limitation

# Examples
[Include 2-3 complete examples of ideal behavior]
"""

# Example 1: Document Search
User: What did the Q3 report say about customer churn?

Thought: The user wants information from a specific document (Q3 report).
I should search their documents for churn-related information.
Action: search_documents
Action Input: {"query": "Q3 report customer churn rate"}

Observation: [1 result] "Q3 2024 Report, page 12: Customer churn decreased
to 2.3% from 2.8% in Q2, driven by improvements in onboarding..."

Thought: I found the relevant information. The Q3 report shows churn
decreased from 2.8% to 2.3%.
Final Answer: According to your Q3 2024 report, customer churn decreased
to 2.3%, down from 2.8% in Q2. The report attributes this improvement to
changes in the onboarding process.

---

# Example 2: Calculation Required
User: If we grew 15% in Q3 and 20% in Q4, what's our total growth?

Thought: I need to calculate compound growth. Two consecutive growth
periods compound multiplicatively, not additively.
Action: calculator
Action Input: {"expression": "((1.15 * 1.20) - 1) * 100"}

Observation: 38.0

Thought: The compound growth is 38%, not 35% (which would be simple addition).
Final Answer: Your total growth across Q3 and Q4 is 38%. This is compound
growth: 1.15 * 1.20 = 1.38, representing a 38% increase from the starting point.

from dataclasses import dataclass
from typing import Callable
import json
import re

@dataclass
class Tool:
    name: str
    description: str
    parameters: dict
    function: Callable

class SimpleAgent:
    def __init__(self, llm, tools: list[Tool], max_iterations: int = 10):
        self.llm = llm
        self.tools = {t.name: t for t in tools}
        self.max_iterations = max_iterations

    def _build_system_prompt(self) -> str:
        tool_desc = "\n".join(
            f"- {t.name}: {t.description}\n  Parameters: {json.dumps(t.parameters)}"
            for t in self.tools.values()
        )
        return f"""You are a helpful assistant with access to tools.

Tools:
{tool_desc}

To use a tool, write:
Thought: [reasoning]
Action: [tool_name]
Action Input: {{"param": "value"}}

When done:
Thought: [summary]
Final Answer: [response]"""

    def run(self, query: str) -> str:
        messages = [
            {"role": "system", "content": self._build_system_prompt()},
            {"role": "user", "content": query}
        ]

        for _ in range(self.max_iterations):
            response = self.llm.chat(messages)
            messages.append({"role": "assistant", "content": response})

            # Check for final answer
            if "Final Answer:" in response:
                return response.split("Final Answer:")[-1].strip()

            # Parse and execute tool call
            action_match = re.search(r"Action:\s*(.+)", response)
            input_match = re.search(r"Action Input:\s*({.+})", response, re.DOTALL)

            if action_match and input_match:
                tool_name = action_match.group(1).strip()
                try:
                    params = json.loads(input_match.group(1))

                    if tool_name in self.tools:
                        result = self.tools[tool_name].function(**params)
                        observation = f"Observation: {result}"
                    else:
                        observation = f"Observation: Error - Unknown tool '{tool_name}'"
                except Exception as e:
                    observation = f"Observation: Error - {str(e)}"

                messages.append({"role": "user", "content": observation})
            else:
                messages.append({
                    "role": "user",
                    "content": "Please use the correct format: Action: tool_name\nAction Input: {params}"
                })

        return "Max iterations reached without final answer."