Files
langgraph-codeact/langgraph_codeact/__init__.py
T
Harrison Chase 68b0ee8bfa cr
2025-03-20 13:54:26 -07:00

85 lines
3.1 KiB
Python

import inspect
from collections import ChainMap
from typing import Any, Callable, Optional, Sequence, Callable, Union
from langchain_core.language_models import BaseChatModel
from langchain_core.messages import AIMessage, BaseMessage, MessageLikeRepresentation
from langchain_core.tools import BaseTool
from langchain_core.tools import tool as create_tool
from langgraph.checkpoint.base import BaseCheckpointSaver
from langgraph.func import entrypoint, task
from langgraph.store.base import BaseStore
from langgraph.graph import StateGraph, MessagesState
from langgraph.types import Command
from typing import Literal, Tuple
class CodeActState(MessagesState):
"""state for codeact agent."""
script: str
context: dict
def create_codeact(
model: BaseChatModel,
tools: Sequence[Union[BaseTool, Callable]],
eval_fn: Callable[[str, dict[str, Callable]], Tuple[str, dict]],
*,
prompt: Optional[str] = None,
):
_tools = [t if isinstance(t, BaseTool) else create_tool(t) for t in tools]
# create the prompt
prompt = f"""
{prompt or ""}
You will be given a task to perform. You should output either
- a Python code snippet that provides the solution to the task, or a step towards the solution. Any output you want to extract from the code should be printed to the console. Code should be output in a fenced code block.
- text to be shown directly to the user, if you want to ask for more information or provide the final answer.
In addition to the Python Standard Library, you can use the following functions:
"""
for tool in _tools:
prompt += f'''
def {tool.name}{str(inspect.signature(tool.func))}:
"""{tool.description}"""
...
'''
prompt += """
Variables defined at the top level of previous code snippets can be referenced in your code.
Reminder: use python code snippets to call tools"""
def call_model(state: CodeActState) -> Command[Literal["__end__", "sandbox"]]:
messages = [{"role": "system", "content": prompt}] + state["messages"]
msg = model.invoke(messages)
if "```" in msg.content:
# get content between fences
code = msg.content.split("```")[1]
# remove first line, which is the language or empty string
code = "\n".join(code.splitlines()[1:])
return Command(goto="sandbox", update={"messages": [msg], "script": code})
else:
# no code block, return None
return Command(goto="__end__", update={"messages": [msg], "script": ""})
def sandbox(state: CodeActState):
script = state["script"]
old_context = state.get("context", {})
context = {**old_context, **{tool.name: tool.func for tool in _tools}}
# execute the script
output, new_vars = eval_fn(script, context)
new_context = {**old_context, **new_vars}
return {"messages": [{"role": "user", "content": output}], "context": new_context}
agent = StateGraph(CodeActState)
agent.add_node(call_model)
agent.add_node(sandbox)
agent.add_edge("__start__", "call_model")
agent.add_edge("sandbox", "call_model")
return agent