Agent工程化实战：从零构建可扩展的AI编程助手

┌─────────────────────────────────────────────────────┐
│                    Orchestrator                      │
│  (规划器 + 执行器 + 反思器)                         │
└────────────┬────────────────────────────────────────┘
             │
     ┌───────┼───────┬──────────────┐
     ▼       ▼       ▼              ▼
┌─────────┐ ┌─────────┐ ┌─────────┐ ┌─────────┐
│ Tool    │ │ Memory  │ │ Logger  │ │ Rate    │
│ Registry│ │ Manager │ │ & Tracer│ │ Limiter │
└─────────┘ └─────────┘ └─────────┘ └─────────┘

# tool_base.py
from pydantic import BaseModel, Field
from typing import Callable, Dict, Any, Optional, Type
import json
import inspect

class ToolSchema(BaseModel):
    name: str
    description: str
    parameters: Dict[str, Any]  # JSON Schema

class Tool:
    def __init__(
        self,
        name: str,
        description: str,
        func: Callable,
        schema: Optional[Dict[str, Any]] = None,
        timeout: float = 30.0
    ):
        self.name = name
        self.description = description
        self.func = func
        self.timeout = timeout
        self.schema = schema or self._infer_schema(func)

    def _infer_schema(self, func):
        sig = inspect.signature(func)
        params = {}
        for p_name, p_param in sig.parameters.items():
            if p_name == "self" or p_name == "kwargs":
                continue
            param_type = "string"
            if p_param.annotation is not inspect.Parameter.empty:
                if p_param.annotation == int:
                    param_type = "integer"
                elif p_param.annotation == float:
                    param_type = "number"
                elif p_param.annotation == bool:
                    param_type = "boolean"
                elif p_param.annotation == list:
                    param_type = "array"
            params[p_name] = {"type": param_type}
        return {
            "type": "object",
            "properties": params,
            "required": list(params.keys())
        }

    async def execute(self, **kwargs) -> Any:
        # 超时控制
        import asyncio
        try:
            return await asyncio.wait_for(
                asyncio.to_thread(self.func, **kwargs),
                timeout=self.timeout
            )
        except asyncio.TimeoutError:
            raise TimeoutError(f"Tool {self.name} execution timeout")

class ToolRegistry:
    def __init__(self):
        self._tools: Dict[str, Tool] = {}

    def register(self, tool: Tool):
        self._tools[tool.name] = tool

    def get_tool(self, name: str) -> Optional[Tool]:
        return self._tools.get(name)

    def list_schemas(self):
        return [
            {
                "type": "function",
                "function": {
                    "name": t.name,
                    "description": t.description,
                    "parameters": t.schema
                }
            }
            for t in self._tools.values()
        ]

# memory.py
import chromadb
from chromadb.utils import embedding_functions
from typing import List, Dict, Any
import json
from collections import deque

class MemoryManager:
    def __init__(self, max_short_term=20, collection_name="agent_memory"):
        self.short_term = deque(maxlen=max_short_term)  # 消息列表
        self.working_memory = {}  # 当前任务状态
        # 长期向量存储
        self.client = chromadb.PersistentClient(path="./chroma_db")
        self.embed_fn = embedding_functions.DefaultEmbeddingFunction()
        self.collection = self.client.get_or_create_collection(
            name=collection_name,
            embedding_function=self.embed_fn
        )

    def add_short_term(self, role: str, content: str):
        self.short_term.append({"role": role, "content": content})

    def get_short_term(self) -> List[Dict]:
        return list(self.short_term)

    def add_long_term(self, text: str, metadata: Dict = None):
        # 生成唯一ID
        import uuid
        doc_id = str(uuid.uuid4())
        self.collection.add(
            documents=[text],
            metadatas=[metadata or {}],
            ids=[doc_id]
        )

    def search_long_term(self, query: str, top_k=3) -> List[str]:
        results = self.collection.query(query_texts=[query], n_results=top_k)
        return results['documents'][0] if results['documents'] else []

    def update_working(self, key: str, value: Any):
        self.working_memory[key] = value

    def get_working(self, key: str, default=None):
        return self.working_memory.get(key, default)

# planner.py
import json
import openai
from typing import List, Dict, Any, Optional
from tool_base import ToolRegistry
from memory import MemoryManager

class ReActPlanner:
    def __init__(
        self,
        llm_client,
        model: str = "gpt-4-turbo",
        max_steps: int = 10,
        tool_registry: Optional[ToolRegistry] = None,
        memory: Optional[MemoryManager] = None
    ):
        self.client = llm_client
        self.model = model
        self.max_steps = max_steps
        self.tools = tool_registry
        self.memory = memory
        self.system_prompt = self._build_system_prompt()

    def _build_system_prompt(self):
        tools_desc = ""
        if self.tools:
            for t in self.tools._tools.values():
                tools_desc += f"- {t.name}: {t.description}\n"
        return f"""You are an AI programming assistant with tool access.
You must respond with a JSON object containing:
- "thought": your reasoning about the current step
- "action": the tool name to call (or "finish" to return final answer)
- "action_input": a dict of parameters for the tool (or final answer string if action is "finish")

Available tools:
{tools_desc}

If you need to finish, action="finish" and action_input is your final answer.
Always think step by step. If a tool fails, try an alternative approach.
"""

    async def step(self, user_query: str) -> Dict:
        # 构建消息上下文
        messages = [
            {"role": "system", "content": self.system_prompt},
        ]
        # 加入短期记忆
        if self.memory:
            short = self.memory.get_short_term()
            messages.extend(short)
        # 加入长期记忆（检索相关）
        if self.memory:
            long_docs = self.memory.search_long_term(user_query, top_k=2)
            if long_docs:
                context = "\n".join(long_docs)
                messages.append({"role": "system", "content": f"Relevant past experiences:\n{context}"})
        # 加入当前用户查询
        messages.append({"role": "user", "content": user_query})

        # 调用LLM，强制JSON输出
        response = self.client.chat.completions.create(
            model=self.model,
            messages=messages,
            response_format={"type": "json_object"},
            temperature=0.2,
        )
        content = response.choices[0].message.content
        try:
            parsed = json.loads(content)
            return parsed
        except json.JSONDecodeError:
            # 容错：尝试提取JSON块
            import re
            match = re.search(r'\{.*\}', content, re.DOTALL)
            if match:
                return json.loads(match.group())
            raise ValueError("LLM did not return valid JSON")

# orchestrator.py
import asyncio
import logging
from typing import Optional
from planner import ReActPlanner
from tool_base import ToolRegistry
from memory import MemoryManager
import time

class Orchestrator:
    def __init__(
        self,
        planner: ReActPlanner,
        tool_registry: ToolRegistry,
        memory: MemoryManager,
        max_retries: int = 2,
    ):
        self.planner = planner
        self.tools = tool_registry
        self.memory = memory
        self.max_retries = max_retries
        self.logger = logging.getLogger("orchestrator")
        self.step_count = 0

    async def run(self, user_query: str) -> str:
        self.step_count = 0
        self.memory.add_short_term("user", user_query)
        # 记录追踪ID
        trace_id = f"trace_{int(time.time()*1000)}"
        self.logger.info({"event": "start", "trace_id": trace_id, "query": user_query})

        # 工作内存初始化
        self.memory.update_working("trace_id", trace_id)
        self.memory.update_working("original_query", user_query)

        final_answer = None
        while self.step_count < self.planner.max_steps:
            self.step_count += 1
            try:
                # 1. 规划
                plan = await self.planner.step(user_query)
                thought = plan.get("thought", "")
                action = plan.get("action", "")
                action_input = plan.get("action_input", {})

                self.logger.info({
                    "event": "plan",
                    "step": self.step_count,
                    "thought": thought,
                    "action": action,
                    "action_input": action_input
                })

                # 2. 判断是否结束
                if action == "finish":
                    final_answer = str(action_input)
                    self.memory.add_short_term("assistant", final_answer)
                    self.logger.info({"event": "finish", "answer": final_answer})
                    break

                # 3. 执行工具
                tool = self.tools.get_tool(action)
                if not tool:
                    observation = f"Error: Tool '{action}' not found."
                else:
                    # 带重试执行
                    for retry in range(self.max_retries + 1):
                        try:
                            result = await tool.execute(**action_input)
                            observation = str(result)
                            break
                        except Exception as e:
                            observation = f"Error (attempt {retry+1}): {str(e)}"
                            if retry == self.max_retries:
                                self.logger.error({"event": "tool_failed", "tool": action, "error": str(e)})
                            await asyncio.sleep(0.5 * (retry+1))  # 退避

                # 4. 将观察写入短期记忆，并更新工作内存
                self.memory.add_short_term("assistant", f"Action: {action}\nObservation: {observation}")
                self.memory.update_working("last_observation", observation)
                self.memory.update_working("last_action", action)

                self.logger.info({
                    "event": "observation",
                    "step": self.step_count,
                    "observation": observation[:200]  # 截断日志
                })

                # 5. 下一轮循环将基于新观察继续规划
                # 将用户查询重置为“继续”，但保留上下文通过short_term传递
                user_query = "Continue with the next step based on the observation."

            except Exception as e:
                self.logger.error({"event": "step_error", "step": self.step_count, "error": str(e)})
                # 错误恢复：向记忆中加入错误，并尝试继续
                self.memory.add_short_term("system", f"Error occurred: {str(e)}. Try an alternative approach.")
                user_query = "Recover from error and try another approach."

        if final_answer is None:
            final_answer = "Maximum steps reached without final answer."
        return final_answer

# tools_impl.py
import os
import subprocess
import glob
from pathlib import Path

def read_file(path: str) -> str:
    """Read the content of a file."""
    with open(path, 'r', encoding='utf-8') as f:
        return f.read()

def write_file(path: str, content: str) -> str:
    """Write content to a file (overwrites if exists)."""
    Path(os.path.dirname(path)).mkdir(parents=True, exist_ok=True)
    with open(path, 'w', encoding='utf-8') as f:
        f.write(content)
    return f"File {path} written successfully."

def execute_shell(command: str, cwd: str = ".") -> str:
    """Execute a shell command and return stdout+stderr."""
    result = subprocess.run(
        command,
        shell=True,
        cwd=cwd,
        capture_output=True,
        text=True,
        timeout=60
    )
    output = result.stdout + result.stderr
    if result.returncode != 0:
        return f"Command failed with code {result.returncode}:\n{output}"
    return output

def search_code(pattern: str, root_dir: str = ".") -> str:
    """Search for a regex pattern in all .py files under root_dir."""
    matches = []
    for filepath in glob.glob(f"{root_dir}/**/*.py", recursive=True):
        try:
            with open(filepath, 'r', encoding='utf-8') as f:
                for i, line in enumerate(f.readlines(), 1):
                    if pattern in line:  # 简单字符串匹配，可改为regex
                        matches.append(f"{filepath}:{i}: {line.strip()[:100]}")
        except Exception:
            continue
    if not matches:
        return "No matches found."
    return "\n".join(matches[:20])  # 限制结果数量

# main.py
import openai
import asyncio
import logging
from tool_base import ToolRegistry, Tool
from tools_impl import read_file, write_file, execute_shell, search_code
from memory import MemoryManager
from planner import ReActPlanner
from orchestrator import Orchestrator

def setup_logging():
    logging.basicConfig(level=logging.INFO)
    # 使用JSON格式的日志（便于采集）
    import json_logging
    json_logging.init_non_web(enable_json=True)

async def main():
    # 1. 初始化LLM客户端（使用OpenAI，或兼容接口）
    client = openai.AsyncOpenAI(
        api_key=os.getenv("OPENAI_API_KEY"),
        base_url=os.getenv("OPENAI_BASE_URL", "https://api.openai.com/v1")
    )

    # 2. 注册工具
    registry = ToolRegistry()
    registry.register(Tool("read_file", "Read content of a file", read_file))
    registry.register(Tool("write_file", "Write content to a file", write_file))
    registry.register(Tool("execute_shell", "Execute a shell command", execute_shell))
    registry.register(Tool("search_code", "Search for a pattern in Python code", search_code))

    # 3. 记忆管理器
    memory = MemoryManager(max_short_term=30)

    # 4. 规划器
    planner = ReActPlanner(
        llm_client=client,
        model="gpt-4-turbo",
        max_steps=8,
        tool_registry=registry,
        memory=memory
    )

    # 5. 编排器
    orchestrator = Orchestrator(planner, registry, memory, max_retries=2)

    # 6. 执行一个实际编程任务
    query = """
    I have a Python project in ./my_project. I need to find all functions that have no docstring,
    and then add a standard docstring to each. The docstring should be:
    \"\"\"Describe what this function does.\"\"\"
    Please don't change any other code. Use search_code to find functions, read_file to read, write_file to modify.
    """
    answer = await orchestrator.run(query)
    print("Final Answer:", answer)

if __name__ == "__main__":
    setup_logging()
    asyncio.run(main())

# rate_limiter.py
import asyncio
import time

class TokenBucket:
    def __init__(self, rate: float, capacity: int):
        self.rate = rate  # tokens per second
        self.capacity = capacity
        self.tokens = capacity
        self.last_refill = time.monotonic()
        self._lock = asyncio.Lock()

    async def acquire(self, tokens=1):
        async with self._lock:
            now = time.monotonic()
            elapsed = now - self.last_refill
            self.tokens = min(self.capacity, self.tokens + elapsed * self.rate)
            self.last_refill = now
            if self.tokens >= tokens:
                self.tokens -= tokens
                return True
            return False

# 在Orchestrator.run中，每次调用LLM和工具前 await rate_limiter.acquire()

# 在Orchestrator.run循环中，每隔几步保存检查点
async def _checkpoint(self):
    import pickle
    state = {
        "short_term": list(self.memory.short_term),
        "working": self.memory.working_memory,
        "step": self.step_count
    }
    with open(f"/tmp/agent_checkpoint_{self.memory.get_working('trace_id')}.pkl", "wb") as f:
        pickle.dump(state, f)

# 在execute_shell中增加安全检查
ALLOWED_COMMANDS = {"ls", "cat", "grep", "find", "python", "pip"}
def execute_shell_safe(command: str, cwd: str = "."):
    cmd_parts = command.split()
    if cmd_parts and cmd_parts[0] not in ALLOWED_COMMANDS:
        raise ValueError(f"Command {cmd_parts[0]} not allowed")
    # 还可以检查路径是否在cwd内
    return execute_shell(command, cwd)

# 修改planner返回结构：{"actions": [{"name": "read_file", "input": {...}}, ...]}
# 在orchestrator中并发执行
async def _execute_parallel(self, actions):
    tasks = []
    for act in actions:
        tool = self.tools.get_tool(act["name"])
        if tool:
            tasks.append(tool.execute(**act["input"]))
        else:
            tasks.append(asyncio.sleep(0, result=f"Error: tool {act['name']} not found"))
    results = await asyncio.gather(*tasks, return_exceptions=True)
    return results

def run_flake8(path: str = ".") -> str:
    result = subprocess.run(
        ["flake8", path, "--format=json"],
        capture_output=True,
        text=True
    )
    if result.stdout:
        import json
        data = json.loads(result.stdout)
        # 格式化输出
        lines = []
        for filepath, errors in data.items():
            for err in errors:
                lines.append(f"{filepath}:{err['line']}:{err['column']} {err['code']} {err['text']}")
        return "\n".join(lines[:30])
    return "No flake8 errors found."

# 注册工具
registry.register(Tool("run_flake8", "Run flake8 and return errors", run_flake8))

# api.py
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import asyncio

app = FastAPI()
class TaskRequest(BaseModel):
    query: str
    trace_id: Optional[str] = None

@app.post("/run")
async def run_task(req: TaskRequest):
    # 每个请求新建Orchestrator实例（或从池中获取）
    orchestrator = create_orchestrator()
    result = await orchestrator.run(req.query)
    return {"result": result, "trace_id": req.trace_id}