V · 6 天前

基于 LlamaIndex 实现 CodeAct Agent:代码执行工作流的技术架构与原理

CodeAct 作为 AI 辅助系统的一种先进范式,实现了自然语言处理与代码执行能力的深度融合。通过构建自定义代码执行代理,开发者能够精确控制应用程序中代码的生成、执行及管理流程。本文将详细阐述如何利用 LlamaIndex 框架从底层构建 CodeAct Agent,深入剖析其内部工作机制,以及如何在预构建解决方案的基础上进行定制化扩展。

CodeAct 技术范式的核心特性

从技术架构角度而言,CodeAct 赋予 AI 助手以下关键能力:

  1. 基于自然语言指令的代码生成:将用户语义需求转换为可执行的程序代码
  2. 安全的代码执行环境:在隔离的受控环境中运行生成的代码
  3. 执行结果分析:获取并解析代码执行的输出与返回值
  4. 迭代式优化:基于执行结果智能调整解决方案

image.png

技术架构与核心组件

CodeAct Agent 的技术架构由以下相互依存的关键组件构成:

  1. 代码执行环境:提供安全隔离的代码运行时环境,确保代码执行不影响宿主系统
  2. 工作流定义系统:规范代码生成、执行与结果处理的逻辑流程
  3. 提示工程机制:引导大语言模型以特定格式生成符合语法与逻辑要求的可执行代码
  4. 状态管理系统:维护对话历史、执行上下文与计算结果的持久化存储

代码实现

1、基础工具函数实现

首先需要定义代理可调用的基础函数集。以下实现了一组基本数学运算工具函数:

 def add(a: int, b: int) -> int:
    """将两个数字相加"""
    return a + b

def subtract(a: int, b: int) -> int:
    """两个数字相减"""
    return a - b

def multiply(a: int, b: int) -> int:
    """两个数字相乘"""
    return a * b

def divide(a: int, b: int) -> float:
    """两个数字相除"""
     return a / b

2、代码执行环境构建

接下来,实现代码执行环境,该环境需具备在多次调用间保持状态的能力。

SimpleCodeExecutor

类实现了这一功能:

 from typing import Any, Dict, Tuple
import io
import contextlib
import ast
import traceback

class SimpleCodeExecutor:
    """
    一个简单的代码执行器,可以在状态持久化的情况下运行Python代码。
    这个执行器在多次执行之间维护全局和局部状态,
    允许变量在多次代码运行中持久化。
    注意:不适合在生产环境中使用!谨慎使用。
    """
    def __init__(self, locals: Dict[str, Any], globals: Dict[str, Any]):
        """
        初始化代码执行器。
        参数:
            locals: 在执行上下文中使用的局部变量
            globals: 在执行上下文中使用的全局变量
        """
        # 在执行之间持久化的状态
        self.globals = globals
        self.locals = locals
    def execute(self, code: str) -> str:
        """
        执行Python代码并捕获输出和返回值。
        参数:
            code: 要执行的Python代码
        返回:
            包含输出和返回值的字符串
        """
        # 捕获标准输出和标准错误
        stdout = io.StringIO()
        stderr = io.StringIO()
        output = ""
        return_value = None
        try:
            # 执行并捕获输出
            with contextlib.redirect_stdout(
                stdout
            ), contextlib.redirect_stderr(stderr):
                # 尝试检测是否有返回值(最后一个表达式)
                try:
                    tree = ast.parse(code)
                    last_node = tree.body[-1] if tree.body else None
                    # 如果最后一个语句是表达式,捕获它的值
                    if isinstance(last_node, ast.Expr):
                        # 分割代码以添加返回值赋值
                        last_line = code.rstrip().split("\n")[-1]
                        exec_code = (
                            code[: -len(last_line)]
                            + "\n__result__ = "
                            + last_line
                        )
                        # 执行修改后的代码
                        exec(exec_code, self.globals, self.locals)
                        return_value = self.locals.get("__result__")
                    else:
                        # 正常执行
                        exec(code, self.globals, self.locals)
                except:
                    # 如果解析失败,按原样执行代码
                    exec(code, self.globals, self.locals)
            # 获取输出
            output = stdout.getvalue()
            if stderr.getvalue():
                output += "\n" + stderr.getvalue()
        except Exception as e:
            # 捕获异常信息
            output = f"Error: {type(e).__name__}: {str(e)}\n"
            output += traceback.format_exc()
        if return_value is not None:
            output += "\n\n" + str(return_value)
         return output

3、工作流事件定义

为了规范化代码执行流程,需要定义控制工作流程的事件类型:

 from llama_index.core.llms import ChatMessage
from llama_index.core.workflow import Event

class InputEvent(Event):
    input: list[ChatMessage]

class StreamEvent(Event):
    delta: str

class CodeExecutionEvent(Event):
     code: str

4、CodeAct Agent 工作流实现

下面是完整的 CodeAct Agent 工作流实现,它将所有组件整合为一个功能完整的系统:

 import inspect
import re
from typing import Any, Callable, List
from llama_index.core.llms import ChatMessage, LLM
from llama_index.core.memory import ChatMemoryBuffer
from llama_index.core.tools.types import BaseTool
from llama_index.core.workflow import (
    Context,
    Workflow,
    StartEvent,
    StopEvent,
    step,
)
from llama_index.llms.openai import OpenAI

CODEACT_SYSTEM_PROMPT = """
你是一个可以执行代码的有用助手。
根据聊天历史,你可以在<execute>...</execute>标签内编写代码来帮助用户解决问题。
在代码中,你可以引用之前使用过的任何变量或函数。
用户还为你提供了一些预定义函数:
{fn_str}
要执行代码,请在<execute>...</execute>标签之间编写代码。
"""

class CodeActAgent(Workflow):
    def __init__(
        self,
        fns: List[Callable],
        code_execute_fn: Callable,
        llm: LLM | None = None,
        **workflow_kwargs: Any,
    ) -> None:
        super().__init__(**workflow_kwargs)
        self.fns = fns or []
        self.code_execute_fn = code_execute_fn
        self.llm = llm or OpenAI(model="gpt-4o-mini")
        # 将函数解析为截断的函数字符串
        self.fn_str = "\n\n".join(
            f'def {fn.__name__}{str(inspect.signature(fn))}:\n    """ {fn.__doc__} """\n    ...'
            for fn in self.fns
        )
        self.system_message = ChatMessage(
            role="system",
            content=CODEACT_SYSTEM_PROMPT.format(fn_str=self.fn_str),
        )
    def _parse_code(self, response: str) -> str | None:
        # 查找<execute>...</execute>标签之间的代码
        matches = re.findall(r"<execute>(.*?)</execute>", response, re.DOTALL)
        if matches:
            return "\n\n".join(matches)
        return None
    @step
    async def prepare_chat_history(
        self, ctx: Context, ev: StartEvent
    ) -> InputEvent:
        # 检查是否设置了内存
        memory = await ctx.get("memory", default=None)
        if not memory:
            memory = ChatMemoryBuffer.from_defaults(llm=self.llm)
        # 获取用户输入
        user_input = ev.get("user_input")
        if user_input is None:
            raise ValueError("user_input kwarg is required")
        user_msg = ChatMessage(role="user", content=user_input)
        memory.put(user_msg)
        # 获取聊天历史
        chat_history = memory.get()
        # 更新上下文
        await ctx.set("memory", memory)
        # 将系统消息添加到聊天历史并返回
        return InputEvent(input=[self.system_message, *chat_history])
    @step
    async def handle_llm_input(
        self, ctx: Context, ev: InputEvent
    ) -> CodeExecutionEvent | StopEvent:
        chat_history = ev.input
        # 流式传输响应
        response_stream = await self.llm.astream_chat(chat_history)
        async for response in response_stream:
            ctx.write_event_to_stream(StreamEvent(delta=response.delta or ""))
        # 保存最终响应,应包含所有内容
        memory = await ctx.get("memory")
        memory.put(response.message)
        await ctx.set("memory", memory)
        # 获取要执行的代码
        code = self._parse_code(response.message.content)
        if not code:
            return StopEvent(result=response)
        else:
            return CodeExecutionEvent(code=code)
    @step
    async def handle_code_execution(
        self, ctx: Context, ev: CodeExecutionEvent
    ) -> InputEvent:
        # 执行代码
        ctx.write_event_to_stream(ev)
        output = self.code_execute_fn(ev.code)
        # 更新内存
        memory = await ctx.get("memory")
        memory.put(ChatMessage(role="assistant", content=output))
        await ctx.set("memory", memory)
        # 获取最新的聊天历史并循环回起点
        chat_history = memory.get()
         return InputEvent(input=[self.system_message, *chat_history])

5、CodeAct Agent 的实例化与调用

完成组件构建后,可以初始化并调用 CodeAct Agent:

 # 使用我们的函数初始化代码执行器
code_executor = SimpleCodeExecutor(
    # 提供访问我们的辅助函数
    locals={
        "add": add,
        "subtract": subtract,
        "multiply": multiply,
        "divide": divide,
    },
    globals={
        # 提供访问所有内置函数
        "__builtins__": __builtins__,
        # 提供访问numpy
        "np": __import__("numpy"),
    },
)

# 创建代理
agent = CodeActAgent(
    fns=[add, subtract, multiply, divide],
    code_execute_fn=code_executor.execute,
    llm=OpenAI(model="gpt-4o-mini", api_key="your_api_key_here"),
)
# 为代理创建上下文
ctx = Context(agent)
# 帮助函数,用于运行代理并输出详细信息
async def run_agent_verbose(agent: CodeActAgent, ctx: Context, query: str):
    handler = agent.run(user_input=query, ctx=ctx)
    print(f"User:  {query}")
    async for event in handler.stream_events():
        if isinstance(event, StreamEvent):
            print(f"{event.delta}", end="", flush=True)
        elif isinstance(event, CodeExecutionEvent):
            print(f"\n-----------\nParsed code:\n{event.code}\n")
    return await handler
# 运行代理与示例
queries = [
    "Calculate the sum of all numbers from 1 to 10",
    "Add 5 and 3, then multiply the result by 2"
]
for query in queries:
    response = await run_agent_verbose(agent, ctx, query)
     print("\n" + "="*50 + "\n")

工作原理分析

CodeAct Agent 的工作流程可分为以下关键阶段:

  1. 用户输入处理:系统接收用户查询,将其添加到对话记忆中,并结合系统提示准备完整的对话上下文。
  2. LLM 代码生成:大语言模型根据对话上下文生成可能包含代码块的响应。系统通过特定标签<execute>...</execute>精确识别可执行代码段。若未检测到代码块,工作流程终止并返回纯文本响应。
  3. 代码执行阶段:系统将识别的代码发送至执行器,在受控环境中运行代码,捕获执行输出和错误信息,并将结果保存至对话记忆。
  4. 响应迭代处理:工作流程循环返回 LLM 处理阶段,使模型能够访问代码执行结果,并根据需要继续生成响应。此迭代过程持续进行,直至生成不包含代码的最终响应。

安全性考量

本文展示的

SimpleCodeExecutor

仅适用于开发环境。生产环境部署时应考虑以下安全措施:

  1. 实现容器化隔离环境(如 Docker),确保代码执行不影响宿主系统
  2. 配置严格的资源限制机制,包括 CPU 使用率、内存上限及最大执行时间
  3. 实施精细的权限控制,严格限制对敏感模块和系统函数的访问
  4. 构建输入验证和安全过滤系统,防止注入攻击

结语

本文详细阐述了基于 LlamaIndex 构建 CodeAct Agent 的完整技术方案,从代码执行环境的构建、工作流事件的定义到完整 Agent 的实现。通过将代码生成与执行能力无缝集成到对话式 AI 系统中,CodeAct Agent 代表了一种新型交互范式,能够将自然语言指令转化为可执行的计算逻辑。

这一技术架构的核心价值在于其可扩展性和灵活性。开发者可以根据特定应用场景定制执行环境、函数库和安全策略,从而构建出专用的智能工具。随着大语言模型能力的不断提升,CodeAct Agent 的应用前景将更加广阔,特别是在以下领域:

  • 专业领域软件开发辅助工具
  • 数据科学探索与可视化系统
  • 教育领域的编程学习平台
  • 企业级自动化工作流构建

未来 CodeAct 技术的发展方向包括更精细的代码生成控制、多语言执行环境支持、更强大的安全隔离机制以及与专业领域知识库的深度集成。这些进步将进一步拓展 AI 辅助编程的边界,使自然语言与代码执行之间的转换更加高效、安全和可靠。

通过深入理解 CodeAct Agent 的核心技术架构,开发者能够构建将大语言模型的自然语言理解能力与代码执行功能有机结合的高度定制化解决方案,为各类应用场景提供强大的智能化支持。

https://avoid.overfit.cn/post/d66bc1e556934aa1b32a3ca58983795c

推荐阅读
关注数
4216
内容数
962
SegmentFault 思否旗下人工智能领域产业媒体,专注技术与产业,一起探索人工智能。
目录
极术微信服务号
关注极术微信号
实时接收点赞提醒和评论通知
安谋科技学堂公众号
关注安谋科技学堂
实时获取安谋科技及 Arm 教学资源
安谋科技招聘公众号
关注安谋科技招聘
实时获取安谋科技中国职位信息