diff --git a/AGENTS.md b/AGENTS.md index 498b4382..26b6c643 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -5,7 +5,7 @@ Core code lives under `src/`: - `src/bub/__main__.py`: Typer CLI entrypoint; `main` (the `bub` script) checks the `app` extra before importing the CLI. -- `src/bub/core/`: agent kernel with no settings or plugin dependencies: agent loop over a tape store (`agent.py`), model protocol and Republic chat model adapter (`model.py`), agent events and their bus (`events.py`), the tape recorder handler (`recorder.py`), system prompt sections (`system_prompt.py`), command prompts (`commands.py`), tools, tape, store, and tape context. +- `src/bub/core/`: agent kernel with no settings or plugin dependencies: agent loop over a tape store (`agent.py`), model protocol and Republic chat model adapter (`model.py`), agent events and their bus (`events.py`), the tape recorder handler (`recorder.py`), system prompt sections (`system_prompt.py`), command prompts (`commands.py`), tools, tape, store, tape context, and code mode (`codemode/`: `CodeMode`, the `run_code` subprocess runner, and the `republic` bridge). - `src/bub/framework.py`: turn orchestration and outbound routing. - `src/bub/hooks/`: hook specifications, execution helpers, and interception contracts. - `src/bub/channels/`: channel bases, routing contracts, admission policy, and lifecycle manager. diff --git a/src/bub/builtin/agent.py b/src/bub/builtin/agent.py index 8f6f6cd7..cee7d20a 100644 --- a/src/bub/builtin/agent.py +++ b/src/bub/builtin/agent.py @@ -229,15 +229,9 @@ def _hook_prompt(self, context: core.PromptContext) -> str: @staticmethod def _code_mode_prompt(context: core.PromptContext) -> str | None: - from bub.builtin.codemode import CODE_TOOLS_STATE_KEY, render_code_mode_prompt, write_tool_stub + from bub.builtin.codemode import CODE_MODE - state = context.state - if (code_tools := state.get(CODE_TOOLS_STATE_KEY)) is None: - return None - stub_path = write_tool_stub( - code_tools, session_id=str(state.get("session_id", "")), workspace=workspace_from_state(state) - ) - return render_code_mode_prompt(stub_path) + return CODE_MODE.prompt(context) def _deferred_tools_prompt(self, context: core.PromptContext) -> str: from bub.builtin.tools import render_deferred_tools_prompt diff --git a/src/bub/builtin/codemode/__init__.py b/src/bub/builtin/codemode/__init__.py index 0a366a67..921bd97e 100644 --- a/src/bub/builtin/codemode/__init__.py +++ b/src/bub/builtin/codemode/__init__.py @@ -1,466 +1,93 @@ -"""Code mode: expose tools to model-written Python through a generated stub and ``run_code``.""" +"""Bub's code mode: :class:`bub.core.codemode.CodeMode` wired to session settings, environments and models. + +``BubAgent`` chooses the code-callable tools of each turn and stores them in the turn +state under ``CODE_TOOLS_STATE_KEY``; the ``code_mode`` command switches it per session. +""" from __future__ import annotations -import asyncio -import contextvars import hashlib -import keyword -import mimetypes -import re -import uuid -from collections.abc import Iterable -from dataclasses import replace +from collections.abc import Iterable, Mapping from pathlib import Path from typing import Any -import republic - -from bub.builtin.codemode.republic_bridge import REPUBLIC_CALL_PREFIX, RepublicSession, decode_value, encode_value from bub.builtin.environment import environment_from_state from bub.builtin.settings import set_session_setting -from bub.core.tools import Tool, ToolContext, ToolExecutor, content_result, model_tools, tool -from bub.environment import CallTool, CodeFailed -from bub.errors import BubError, ErrorKind -from bub.prompt import UserContent +from bub.core.codemode import ( + DEFAULT_RUN_CODE_TIMEOUT_SECONDS, + MEDIA_TOOLS, + RUN_CODE_TOOL_NAME, + CodeMode, + attach_audio, + attach_image, + attach_video, + render_code_mode_prompt, + render_tool_stub, + write_stub, +) +from bub.core.tools import REGISTRY, Tool, ToolContext, tool +from bub.turn import TurnState +from bub.utils import workspace_from_state + +__all__ = [ + "CODE_MODE", + "CODE_MODE_STATE_KEY", + "CODE_TOOLS_STATE_KEY", + "DEFAULT_RUN_CODE_TIMEOUT_SECONDS", + "RUN_CODE_TOOL_NAME", + "attach_audio", + "attach_image", + "attach_video", + "render_code_mode_prompt", + "render_tool_stub", + "run_code", + "set_code_mode", + "write_tool_stub", +] -RUN_CODE_TOOL_NAME = "run_code" CODE_MODE_STATE_KEY = "code_mode" CODE_TOOLS_STATE_KEY = "_runtime_code_tools" -DEFAULT_RUN_CODE_TIMEOUT_SECONDS = 120 - -_RUN_CODE_MEDIA: contextvars.ContextVar[list[UserContent] | None] = contextvars.ContextVar( - "bub_run_code_media", default=None -) - -_STUB_HEADER = '''"""Bub tools available inside `run_code` as `tools.(...)`. - -Every tool is an async function: `await` it (top-level `await` is allowed) and pass keyword -arguments only. Each call returns the structured value described by its return type and -raises an exception when the tool fails. Use `asyncio.gather` to run independent calls -concurrently. - -The `republic` module is also available (preset as a global and importable); it is documented -at the end of this file. -""" - -from typing import Any, Literal, NotRequired, TypedDict - -import republic''' - -_REPUBLIC_STUB = """# --------------------------------------------------------------------------- -# `republic`: media and models inside `run_code`. Models are served by Bub with its -# credentials and settings; every model method is async. Only the members below exist. -# -# class republic.Image / republic.Audio / republic.Video: -# __init__(media_type: str, data: bytes | None = None, url: str | None = None) -# media_type: str; data: bytes | None; url: str | None -# base64_data: str; data_url: str -# `tools.attach_image/audio/video` also take a file path in the environment, an http(s)/gs URL -# or a data URL instead, and load it for you. -# -# republic.system(text: str) -> Message -# republic.user(*content: str | Image | Audio | Video) -> Message -# republic.assistant(*content: str | Image | Audio | Video) -> Message -# Message has .role, .content (list of text and media) and .text. -# -# republic.get_model(spec: str | None = None) -> ChatModel -# A "provider:model" spec; None is the model running this session. -# await model.chat(prompt: str | Message | list[str | Message], **options) -> response -# options: max_tokens, temperature, top_p, top_k, stop, seed, reasoning_effort, extra_body... -# (no tools). A plain string is a user message. -# response.text, .reasoning, .refusal, .finish_reason, .model, -# .token_usage.input_tokens / .output_tokens, .message (append it to continue the conversation) -# model.stream(prompt, **options) -> Stream # same arguments as chat(); not awaited -# async with model.stream(...) as stream: -# async for event in stream: ... -# Events from republic.events: TextDelta(.chunk), ReasoningDelta(.chunk), RefusalDelta(.chunk), -# UsageDelta(.usage), ImageReady(.image), CitationAdded(.citation), Completed(.response, last). -# After iterating: stream.response (as returned by chat()), stream.text, .reasoning, .token_usage. -# republic.get_embedding_model(spec: str) -> EmbeddingModel -# await model.embed(text: str, *, dimensions: int | None = None) -> response # response.vector -# await model.embed_many(texts: list[str], *, dimensions: int | None = None) -> response # response.vectors -# republic.get_decision_model(spec: str) -> DecisionModel -# await model.decide(state, *, questions: dict[str, Question]) -> response -# Questions from republic.decisions: Noul(instructions, criteria=None) (yes/no), -# Choice(instructions, criteria) (pick one), Score(instructions, criteria) (ordered levels). -# response.answers[id] or response.: .noul | .choice/.probabilities/.confidence | .score/... -# -# Example: -# model = republic.get_model() -# chart = republic.Image("image/png", data=open("chart.png", "rb").read()) -# reply = await model.chat([republic.user("Describe this chart.", chart)]) -# print(reply.text) -# async with model.stream("Write a haiku.") as stream: -# async for event in stream: -# if isinstance(event, republic.events.TextDelta): -# print(event.chunk, end="")""" -_JSON_TYPES = {"string": "str", "integer": "int", "number": "float", "boolean": "bool", "null": "None"} - - -def _identifier(name: str) -> str: - ident = re.sub(r"\W", "_", name) - if not ident or ident[0].isdigit() or keyword.iskeyword(ident): - ident = f"_{ident}" - return ident - - -def _is_identifier(name: str) -> bool: - return name.isidentifier() and not keyword.iskeyword(name) - - -def _pascal(name: str) -> str: - return "".join(part[:1].upper() + part[1:] for part in re.split(r"[^0-9A-Za-z]+", name)) - - -_MEDIA_FIELDS = {"media_type", "data", "url"} - - -def _media_class(schema: dict[str, Any]) -> str | None: - """Recognize the JSON schema of ``republic.Image``, ``Audio`` and ``Video`` parameters.""" - title = schema.get("title") - properties = schema.get("properties") - if title in ("Image", "Audio", "Video") and isinstance(properties, dict) and set(properties) == _MEDIA_FIELDS: - return str(title) - return None - - -def _docstring(text: str, indent: str = " ") -> list[str]: - text = text.strip().replace("\\", "\\\\").replace('"""', '\\"\\"\\"') - if not text: - return [] - lines = text.splitlines() - if len(lines) == 1: - return [f'{indent}"""{lines[0]}"""'] - return [f'{indent}"""{lines[0]}', *(f"{indent}{line}".rstrip() for line in lines[1:]), f'{indent}"""'] - - -class _StubBuilder: - """Render tool JSON schemas as Python stub declarations.""" - - def __init__(self) -> None: - self.classes: dict[str, str] = {} - self._refs: dict[tuple[int, str], str] = {} - - def function(self, tool_item: Tool) -> str: - name = _identifier(tool_item.name) - schema = tool_item.parameters or {} - defs = self._defs(schema) - properties = schema.get("properties") or {} - required = set(schema.get("required") or []) - params: list[str] = [] - arg_docs: list[str] = [] - for key, prop in properties.items(): - hint = self.type_hint(prop, defs, f"{_pascal(name)}{_pascal(key)}") - if key in required: - params.append(f"{key}: {hint}") - elif isinstance(prop, dict) and "default" in prop: - params.append(f"{key}: {hint} = {prop['default']!r}") - else: - params.append(f"{key}: {hint} = ...") - description = prop.get("description") if isinstance(prop, dict) else None - if description: - arg_docs.append(f" {key}: {' '.join(description.split())}") - if not all(_is_identifier(key) for key in properties): - signature = "**kwargs: Any" - elif params: - signature = f"*, {', '.join(params)}" - else: - signature = "" - - output = tool_item.output_schema - returns = "Any" if output is None else self.type_hint(output, self._defs(output), f"{_pascal(name)}Result") - doc = tool_item.description.strip() - if arg_docs: - doc = f"{doc}\n\nArgs:\n" + "\n".join(arg_docs) if doc else "Args:\n" + "\n".join(arg_docs) - body = _docstring(doc) or [" ..."] - return "\n".join([f"async def {name}({signature}) -> {returns}:", *body]) - - def type_hint( # noqa: C901 - self, schema: Any, defs: dict[str, Any], name_hint: str, ref_key: tuple[int, str] | None = None - ) -> str: - if not isinstance(schema, dict) or not schema: - return "Any" - if (media := _media_class(schema)) is not None: - return f"republic.{media}" - if local_defs := self._defs(schema): - # Schemas of single parameters carry their own definitions. - defs = {**defs, **local_defs} - if isinstance(ref := schema.get("$ref"), str): - return self._ref(ref, defs) - for key in ("anyOf", "oneOf"): - if isinstance(options := schema.get(key), list): - return self._union(self.type_hint(option, defs, name_hint) for option in options) - if isinstance(all_of := schema.get("allOf"), list) and len(all_of) == 1: - return self.type_hint(all_of[0], defs, name_hint, ref_key) - if "const" in schema: - return f"Literal[{schema['const']!r}]" - if isinstance(enum := schema.get("enum"), list) and enum: - return f"Literal[{', '.join(repr(value) for value in enum)}]" - - schema_type = schema.get("type") - if isinstance(schema_type, list): - return self._union(self.type_hint({**schema, "type": item}, defs, name_hint) for item in schema_type) - if schema_type == "array": - if isinstance(prefix := schema.get("prefixItems"), list): - return f"tuple[{', '.join(self.type_hint(item, defs, name_hint) for item in prefix)}]" - return f"list[{self.type_hint(schema.get('items'), defs, f'{name_hint}Item')}]" - if schema_type == "object": - if isinstance(properties := schema.get("properties"), dict) and properties: - return self._typed_dict(schema, properties, defs, schema.get("title") or name_hint, ref_key) - extra = schema.get("additionalProperties") - if isinstance(extra, dict) and extra: - return f"dict[str, {self.type_hint(extra, defs, f'{name_hint}Value')}]" - return "dict[str, Any]" - return _JSON_TYPES.get(schema_type, "Any") if isinstance(schema_type, str) else "Any" - @staticmethod - def _defs(schema: dict[str, Any]) -> dict[str, Any]: - defs = schema.get("$defs", schema.get("definitions")) - return defs if isinstance(defs, dict) else {} - @staticmethod - def _union(hints: Iterable[str]) -> str: - unique = list(dict.fromkeys(hints)) - return "Any" if "Any" in unique else " | ".join(unique) +def _resolve_model(spec: str | None, context: ToolContext) -> tuple[str, Mapping[str, Any]]: + """Resolve a spec, or the session's model, with Bub's settings: fallbacks, credentials and API bases.""" + agent = context.state.get("_runtime_agent") + settings = getattr(agent, "settings", None) + if settings is None: + raise ValueError("Republic models need a running Bub agent.") + candidate = settings.model_candidates(spec or context.state.get("model") or settings.model)[0] + client_kwargs = settings.model_client_kwargs(candidate.provider_name or candidate.provider) + return f"{candidate.provider}:{candidate.model_id}", client_kwargs - def _ref(self, ref: str, defs: dict[str, Any]) -> str: - name = ref.rsplit("/", 1)[-1] - key = (id(defs), name) - if key in self._refs: - return self._refs[key] - target = defs.get(name) - if not isinstance(target, dict): - return "Any" - hint = self.type_hint(target, defs, name, ref_key=key) - self._refs[key] = hint - return hint - def _typed_dict( - self, - schema: dict[str, Any], - properties: dict[str, Any], - defs: dict[str, Any], - title: str, - ref_key: tuple[int, str] | None, - ) -> str: - base = _identifier(_pascal(title) or "Result") - class_name = base - index = 2 - while class_name in self.classes: - class_name = f"{base}{index}" - index += 1 - # Reserve the name first so recursive references resolve to this class. - self.classes[class_name] = "" - if ref_key is not None: - self._refs[ref_key] = class_name - - required = set(schema.get("required") or []) - fields: list[tuple[str, str, str | None]] = [] - for key, prop in properties.items(): - hint = self.type_hint(prop, defs, f"{class_name}{_pascal(key)}") - if key not in required: - hint = f"NotRequired[{hint}]" - description = prop.get("description") if isinstance(prop, dict) else None - fields.append((key, hint, description)) - source = self._class_source(class_name, fields, schema.get("description")) - - # Reuse an identical class that was declared under the same title by another tool. - for existing, existing_source in self.classes.items(): - if ( - existing != class_name - and re.fullmatch(rf"{re.escape(base)}\d*", existing) - and existing_source == source.replace(class_name, existing) - ): - del self.classes[class_name] - if ref_key is not None: - self._refs[ref_key] = existing - return existing - self.classes[class_name] = source - return class_name +def _stub_dir(session_id: str, workspace: Path) -> Path: + import bub - @staticmethod - def _class_source(name: str, fields: list[tuple[str, str, str | None]], description: object) -> str: - if not all(_is_identifier(key) for key, _, _ in fields): - items = ", ".join(f"{key!r}: {hint}" for key, hint, _ in fields) - return f"{name} = TypedDict({name!r}, {{{items}}})" - lines = [f"class {name}(TypedDict):"] - if isinstance(description, str): - lines.extend(_docstring(description)) - for key, hint, field_description in fields: - if field_description: - lines.append(f" # {' '.join(field_description.split())}") - lines.append(f" {key}: {hint}") - return "\n".join(lines) + return bub.home / "codemode" / hashlib.sha256(f"{workspace}\0{session_id}".encode()).hexdigest()[:16] -def render_tool_stub(tools: Iterable[Tool]) -> str: - """Render a Python stub declaring one function per model-facing tool.""" - builder = _StubBuilder() - functions = [builder.function(tool_item) for tool_item in model_tools(tools)] - return "\n\n\n".join([_STUB_HEADER, *builder.classes.values(), *functions, _REPUBLIC_STUB]) + "\n" +def _state_stub_dir(state: TurnState) -> Path: + return _stub_dir(str(state.get("session_id", "")), workspace_from_state(state)) def write_tool_stub(tools: Iterable[Tool], *, session_id: str, workspace: Path) -> Path: """Write the tool stub under Bub home; the path stays stable while the session's tool set is unchanged.""" - import bub - - content = render_tool_stub(tools) - session_key = hashlib.sha256(f"{workspace}\0{session_id}".encode()).hexdigest()[:16] - content_key = hashlib.sha256(content.encode()).hexdigest()[:12] - path = bub.home / "codemode" / session_key / f"tools-{content_key}.pyi" - if not path.exists(): - path.parent.mkdir(parents=True, exist_ok=True) - temp_path = path.with_name(f".{path.name}.{uuid.uuid4().hex}.tmp") - temp_path.write_text(content, encoding="utf-8") - temp_path.replace(path) - return path - - -def render_code_mode_prompt(stub_path: Path) -> str: - return ( - "\n" - f"More tools are available as async Python functions `tools.(...)` inside `{RUN_CODE_TOOL_NAME}`. " - f"Their signatures, result types and documentation are in the stub file: {stub_path}\n" - "Read the stub before calling a tool you have not used yet. Always `await` tool calls (top-level `await` " - "is allowed) and pass keyword arguments; they return structured values and raise on failure. " - f"`{RUN_CODE_TOOL_NAME}` returns only what the code prints, so print the results you need, and combine " - "several tool calls in one run when possible. To look at an image, audio or video, pass its path, URL " - "or `republic` media object to `tools.attach_image(image=...)`, `tools.attach_audio(audio=...)` or " - f"`tools.attach_video(video=...)`: it is attached to the `{RUN_CODE_TOOL_NAME}` result.\n" - "The code can also use the `republic` module (see the end of the stub) to call models with " - "`republic.get_model()`, " - "`republic.get_embedding_model(spec)` and `republic.get_decision_model(spec)`; `get_model()` without a " - "spec is the model running this session.\n" - "" - ) - - -def _code_tool_caller( - tools_by_name: dict[str, Tool], context: ToolContext, republic_session: RepublicSession -) -> CallTool: - """Serve the code's ``tools.`` and ``republic`` calls.""" - # Code consumes structured results, so this executor does not render them to text. - executor = ToolExecutor(context.events, render=False) - code_context = replace(context, code_mode=True) - - async def call_tool(name: str, arguments: dict[str, Any]) -> Any: - arguments = decode_value(arguments) - if name.startswith(REPUBLIC_CALL_PREFIX): - return encode_value(await republic_session.call(name, arguments)) - tool_item = tools_by_name.get(name) - if tool_item is None: - raise BubError(ErrorKind.INVALID_INPUT, f"unknown tool: tools.{name}") - return encode_value(await executor.execute(tool_item, arguments, context=code_context)) - - return call_tool - - -@tool(name=RUN_CODE_TOOL_NAME, context=True, exposure="direct") -async def run_code( - code: str, timeout_seconds: int = DEFAULT_RUN_CODE_TIMEOUT_SECONDS, *, context: ToolContext -) -> str | dict[str, Any]: - """Run Python code in the environment and return everything it prints. - - Tools are async functions available as `tools.(...)`: await them with keyword arguments - (top-level `await` is allowed). See the tool stub file referenced in the system prompt for their - signatures and result types. The `republic` module provides media loaders and models. Media passed - to `tools.attach_image/audio/video`, as an object, path or URL, is attached to the result. The code is stopped after timeout_seconds. - """ - code_tools = context.state.get(CODE_TOOLS_STATE_KEY) - if code_tools is None: - raise BubError(ErrorKind.INVALID_INPUT, "Code mode is not enabled for this run.") - tools_by_name = {_identifier(item.name): item for item in code_tools} - republic_session = RepublicSession(replace(context, code_mode=True)) - call_tool = _code_tool_caller(tools_by_name, context, republic_session) - - output: list[str] = [] - media: list[UserContent] = [] - token = _RUN_CODE_MEDIA.set(media) - try: - async with asyncio.timeout(timeout_seconds) as deadline: - await environment_from_state(context.state).run_code( - code, tools=list(tools_by_name), call_tool=call_tool, write=output.append - ) - except TimeoutError: - if not deadline.expired(): - raise - raise BubError( - ErrorKind.TOOL, f"Code timed out after {timeout_seconds} seconds", details={"output": "".join(output)} - ) from None - except CodeFailed as exc: - raise BubError( - ErrorKind.TOOL, f"Code raised {exc.error}", details={"output": "".join(output), "traceback": exc.traceback} - ) from exc - except BubError as exc: - raise BubError(exc.kind, exc.message, details={"output": "".join(output), **(exc.details or {})}) from exc - finally: - _RUN_CODE_MEDIA.reset(token) - await republic_session.aclose() - if media: - return content_result(["".join(output), *media] if output else media) - return "".join(output) + return write_stub(render_tool_stub(tools), _stub_dir(session_id, workspace)) -_REMOTE_MEDIA_PREFIXES = ("http://", "https://", "gs://", "data:") -_MEDIA_CLASSES: dict[str, type[republic.Image | republic.Audio | republic.Video]] = { - "image": republic.Image, - "audio": republic.Audio, - "video": republic.Video, -} - - -async def _load_media(kind: str, source: str, context: ToolContext) -> republic.Image | republic.Audio | republic.Video: - """Load media from a URL, a data URL, or a file path in the session environment.""" - if source.startswith(_REMOTE_MEDIA_PREFIXES): - media: republic.Image | republic.Audio | republic.Video = getattr(republic, kind)(source) - return media - media_type, _ = mimetypes.guess_type(source) - if media_type is None: - raise ValueError(f"cannot guess the media type of {source!r}") - environment = environment_from_state(context.state) - data = await environment.read_bytes(environment.resolve_path(source)) - return _MEDIA_CLASSES[kind](media_type, data=data) - - -async def _attach(kind: str, item: republic.Image | republic.Audio | republic.Video | str, context: ToolContext) -> str: - media = _RUN_CODE_MEDIA.get() - if media is None: - raise BubError(ErrorKind.INVALID_INPUT, f"tools.attach_{kind} can only be called from {RUN_CODE_TOOL_NAME}.") - if isinstance(item, str): - try: - item = await _load_media(kind, item, context) - except (OSError, ValueError) as exc: - raise BubError(ErrorKind.INVALID_INPUT, f"Cannot load the {kind}: {exc}") from exc - media.append(item) - return f"The {kind} is attached to the {RUN_CODE_TOOL_NAME} result." - - -@tool(name="attach_image", context=True, exposure="code") -async def attach_image(image: republic.Image | str, *, context: ToolContext) -> str: - """Attach an image to the result of this `run_code` call, so you can see it. - - `image` is a `republic.Image`, a file path in the environment, an http(s) URL, or a data URL. - """ - return await _attach("image", image, context) - - -@tool(name="attach_audio", context=True, exposure="code") -async def attach_audio(audio: republic.Audio | str, *, context: ToolContext) -> str: - """Attach audio to the result of this `run_code` call, so you can hear it. - - `audio` is a `republic.Audio`, a file path in the environment, an http(s) URL, or a data URL. - """ - return await _attach("audio", audio, context) - - -@tool(name="attach_video", context=True, exposure="code") -async def attach_video(video: republic.Video | str, *, context: ToolContext) -> str: - """Attach a video to the result of this `run_code` call, so you can watch it. +CODE_MODE = CodeMode( + lambda state: state.get(CODE_TOOLS_STATE_KEY), + environment=environment_from_state, + models=_resolve_model, + stub_dir=_state_stub_dir, + # The media tools are registered below, so turns select them like any other tool. + media_tools=False, +) +"""Code mode of Bub's agents.""" - `video` is a `republic.Video`, a file path in the environment, an http(s) URL, or a data URL. - """ - return await _attach("video", video, context) +run_code = CODE_MODE.run_code +for _tool in (run_code, *MEDIA_TOOLS): + REGISTRY[_tool.name] = _tool @tool(name="code_mode", context=True, exposure="command") diff --git a/src/bub/builtin/environment.py b/src/bub/builtin/environment.py index 91a4d3f1..f3399dca 100644 --- a/src/bub/builtin/environment.py +++ b/src/bub/builtin/environment.py @@ -1,139 +1,16 @@ -"""Host implementation of :class:`bub.environment.Environment`, provided by Bub's builtin hooks.""" +"""Bub's builtin environment provider: the session's environment from turn state.""" from __future__ import annotations -import asyncio -import contextlib -import os -import shutil -import signal -import sys -from collections.abc import Callable, Mapping, Sequence from pathlib import Path -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING -from bub.environment import ENVIRONMENT_STATE_KEY, CallTool, Environment, Process +from bub.environment import ENVIRONMENT_STATE_KEY, Environment, LocalEnvironment if TYPE_CHECKING: from bub.turn import TurnState -class LocalProcess(Process): - """A host process in its own session, so signals reach the whole process group.""" - - def __init__(self, process: asyncio.subprocess.Process) -> None: - self.process = process - - @property - def pid(self) -> int: - return self.process.pid - - @property - def stdout(self) -> asyncio.StreamReader: - assert self.process.stdout is not None # noqa: S101 - return self.process.stdout - - @property - def stderr(self) -> asyncio.StreamReader: - assert self.process.stderr is not None # noqa: S101 - return self.process.stderr - - @property - def returncode(self) -> int | None: - return self.process.returncode - - async def wait(self) -> int: - return await self.process.wait() - - async def write_stdin(self, data: bytes) -> None: - assert self.process.stdin is not None # noqa: S101 - self.process.stdin.write(data) - await self.process.stdin.drain() - - def close_stdin(self) -> None: - if self.process.stdin is not None: - self.process.stdin.close() - - def signal(self, *, kill: bool) -> None: - with contextlib.suppress(ProcessLookupError): - if os.name != "nt": - os.killpg(self.process.pid, signal.SIGKILL if kill else signal.SIGTERM) - elif self.returncode is None: - if kill: - self.process.kill() - else: - self.process.terminate() - - def is_running(self) -> bool: - if os.name == "nt": - return self.returncode is None - try: - os.killpg(self.process.pid, 0) - except ProcessLookupError: - return False - except PermissionError: - # EPERM does not establish that the group is gone (in particular - # while its leader is exiting on macOS). - return True - return True - - -class LocalEnvironment(Environment): - """Run processes and access files directly on the host.""" - - SHELL = shutil.which("bash") or shutil.which("sh") if os.name != "nt" else None - - def __init__(self, workspace: str | Path | None = None) -> None: - self.workspace = str(workspace) if workspace is not None else None - self.python = sys.executable - - async def spawn( - self, command: str | Sequence[str], *, cwd: str | None = None, env: Mapping[str, str] | None = None - ) -> LocalProcess: - options: dict[str, Any] = { - "cwd": cwd or self.workspace, - "env": {**os.environ, **env} if env else None, - "stdin": asyncio.subprocess.PIPE, - "stdout": asyncio.subprocess.PIPE, - "stderr": asyncio.subprocess.PIPE, - "start_new_session": os.name != "nt", - } - if isinstance(command, str): - process = await asyncio.create_subprocess_shell(command, executable=self.SHELL, **options) - else: - process = await asyncio.create_subprocess_exec(*command, **options) - return LocalProcess(process) - - async def read_text(self, path: str) -> str: - return await asyncio.to_thread(Path(path).read_text, encoding="utf-8") - - async def read_bytes(self, path: str) -> bytes: - return await asyncio.to_thread(Path(path).read_bytes) - - async def write_text(self, path: str, content: str) -> None: - def write() -> None: - target = Path(path) - target.parent.mkdir(parents=True, exist_ok=True) - target.write_text(content, encoding="utf-8") - - await asyncio.to_thread(write) - - async def run_code( - self, code: str, *, tools: Sequence[str], call_tool: CallTool, write: Callable[[str], None] - ) -> None: - from bub.builtin.codemode.code_runner import run_code_in_subprocess - - await run_code_in_subprocess(self, code, tools=tools, call_tool=call_tool, write=write) - - def resolve_path(self, path: str) -> str: - expanded = Path(path).expanduser() - if expanded.is_absolute(): - return str(expanded) - if self.workspace is None: - raise ValueError(f"relative path '{path}' is not allowed without a workspace") - return str((Path(self.workspace) / expanded).resolve()) - - def environment_from_state(state: TurnState) -> Environment: """Return the session's environment, or a host environment for the state's workspace when none was provided. diff --git a/src/bub/builtin/hook_impl.py b/src/bub/builtin/hook_impl.py index cfc302b2..4454a455 100644 --- a/src/bub/builtin/hook_impl.py +++ b/src/bub/builtin/hook_impl.py @@ -288,7 +288,7 @@ async def provide_lifespan(self) -> AsyncIterator[None]: @hookimpl def provide_environment(self, session_id: str, workspace: Path) -> Environment: - from bub.builtin.environment import LocalEnvironment + from bub.environment import LocalEnvironment return LocalEnvironment(workspace) diff --git a/src/bub/builtin/shell_manager.py b/src/bub/builtin/shell_manager.py index 9cebde9d..e6bea2a8 100644 --- a/src/bub/builtin/shell_manager.py +++ b/src/bub/builtin/shell_manager.py @@ -7,8 +7,7 @@ from contextvars import ContextVar from dataclasses import dataclass, field -from bub.builtin.environment import LocalEnvironment -from bub.environment import Environment, Process +from bub.environment import Environment, LocalEnvironment, Process @dataclass(eq=False) diff --git a/src/bub/core/codemode/__init__.py b/src/bub/core/codemode/__init__.py new file mode 100644 index 00000000..50062a5c --- /dev/null +++ b/src/bub/core/codemode/__init__.py @@ -0,0 +1,306 @@ +"""Code mode: let the model call tools from Python it writes, through one ``run_code`` tool. + +:class:`CodeMode` bundles what an agent needs: the ``run_code`` tool, a system prompt +section describing the code-callable tools as a Python stub, and the ``attach_*`` media +tools. It works with a bare :class:`bub.core.Agent`:: + + code_mode = CodeMode([search, fetch], models="openai:gpt-5") + agent = Agent(model=model, tools=[code_mode.run_code], system_prompt=["You are ...", code_mode.prompt]) + +Code runs in an :class:`bub.environment.Environment`, a Python process on the host by +default. Calls the code makes go through the running agent's tool events, so hooks and +handlers see them with ``context.code_mode`` set. +""" + +from __future__ import annotations + +import asyncio +import contextvars +import hashlib +import mimetypes +import uuid +from collections.abc import Callable, Iterable, Mapping +from dataclasses import dataclass, replace +from pathlib import Path +from typing import Any + +import republic + +from bub.core.codemode.republic_bridge import ( + REPUBLIC_CALL_PREFIX, + ModelResolver, + RepublicSession, + decode_value, + default_model_resolver, + encode_value, +) +from bub.core.codemode.stub import _identifier, render_tool_stub +from bub.core.system_prompt import PromptContext +from bub.core.tools import Tool, ToolContext, ToolExecutor, content_result, model_tools +from bub.environment import CallTool, CodeFailed, Environment, LocalEnvironment +from bub.errors import BubError, ErrorKind +from bub.prompt import UserContent +from bub.turn import TurnState + +__all__ = [ + "DEFAULT_RUN_CODE_TIMEOUT_SECONDS", + "MEDIA_TOOLS", + "RUN_CODE_TOOL_NAME", + "CodeMode", + "ModelResolver", + "attach_audio", + "attach_image", + "attach_video", + "default_model_resolver", + "render_code_mode_prompt", + "render_tool_stub", + "write_stub", +] + +RUN_CODE_TOOL_NAME = "run_code" +DEFAULT_RUN_CODE_TIMEOUT_SECONDS = 120 + +CodeTools = Iterable[Tool] | Callable[[TurnState], Iterable[Tool] | None] +"""The tools callable from code: fixed, or chosen per run from the turn state (None disables ``run_code``).""" + + +@dataclass(frozen=True) +class _CodeRun: + environment: Environment + media: list[UserContent] + + +_ACTIVE_RUN: contextvars.ContextVar[_CodeRun | None] = contextvars.ContextVar("bub_run_code", default=None) + + +class CodeMode: + """The ``run_code`` tool and the system prompt section for a set of code-callable tools. + + Args: + tools: Tools callable from code as ``tools.(...)``. A callable receives the + turn state and returns the tools for that run, or None when code mode is off. + environment: Where the code runs, or a callable returning it for the turn state. + Defaults to a :class:`bub.environment.LocalEnvironment` in the working directory. + models: What ``republic.get_*model(spec)`` in the code resolves to: a default + ``provider:model`` spec for ``get_model()``, or a :data:`ModelResolver`. By default + specs must be explicit and clients read credentials from the environment. + stub_dir: Write the tool stub to a file in this directory (or the one returned for + the turn state) and point the prompt at it, for agents that can read files. + None puts the stub in the system prompt itself. + media_tools: Also offer ``attach_image``, ``attach_audio`` and ``attach_video`` to + the code, which attach media to the ``run_code`` result. + name: Name of the ``run_code`` tool. + """ + + def __init__( + self, + tools: CodeTools = (), + *, + environment: Environment | Callable[[TurnState], Environment] | None = None, + models: str | ModelResolver | None = None, + stub_dir: Path | Callable[[TurnState], Path] | None = None, + media_tools: bool = True, + name: str = RUN_CODE_TOOL_NAME, + ) -> None: + self._tools = tools if callable(tools) else list(tools) + self._environment = LocalEnvironment() if environment is None else environment + self.resolve_model = models if callable(models) else default_model_resolver(models) + self.stub_dir = stub_dir + self.media_tools = media_tools + self.run_code: Tool = Tool.from_callable( + self._run_code, name=name, description=_RUN_CODE_DESCRIPTION, context=True, exposure="direct" + ) + """The model-facing tool; add it to the agent's tools.""" + + def code_tools(self, state: TurnState) -> list[Tool] | None: + """Return the model-safe tools callable from code in this run, or None when code mode is off.""" + tools = self._tools(state) if callable(self._tools) else self._tools + if tools is None: + return None + result = model_tools(tools) + if self.media_tools: + names = {tool_item.name for tool_item in result} + result.extend(tool_item for tool_item in MEDIA_TOOLS if tool_item.name not in names) + return result + + def environment(self, state: TurnState) -> Environment: + """Return the environment the code runs in for this run.""" + return self._environment(state) if callable(self._environment) else self._environment + + def prompt(self, context: PromptContext) -> str | None: + """System prompt section describing the code-callable tools; add it to the agent's system prompt.""" + state = context.state + if (tools := self.code_tools(state)) is None: + return None + stub = render_tool_stub(tools) + if self.stub_dir is not None: + directory = self.stub_dir(state) if callable(self.stub_dir) else self.stub_dir + return render_code_mode_prompt(write_stub(stub, directory), tool_name=self.run_code.name) + return render_code_mode_prompt(stub, tool_name=self.run_code.name) + + async def _run_code( + self, code: str, timeout_seconds: int = DEFAULT_RUN_CODE_TIMEOUT_SECONDS, *, context: ToolContext + ) -> str | dict[str, Any]: + code_tools = self.code_tools(context.state) + if code_tools is None: + raise BubError(ErrorKind.INVALID_INPUT, "Code mode is not enabled for this run.") + tools_by_name = {_identifier(item.name): item for item in code_tools} + code_context = replace(context, code_mode=True) + republic_session = RepublicSession(code_context, self.resolve_model) + call_tool = _code_tool_caller(tools_by_name, code_context, republic_session) + environment = self.environment(context.state) + + output: list[str] = [] + media: list[UserContent] = [] + token = _ACTIVE_RUN.set(_CodeRun(environment, media)) + try: + async with asyncio.timeout(timeout_seconds) as deadline: + await environment.run_code(code, tools=list(tools_by_name), call_tool=call_tool, write=output.append) + except TimeoutError: + if not deadline.expired(): + raise + raise BubError( + ErrorKind.TOOL, f"Code timed out after {timeout_seconds} seconds", details={"output": "".join(output)} + ) from None + except CodeFailed as exc: + raise BubError( + ErrorKind.TOOL, + f"Code raised {exc.error}", + details={"output": "".join(output), "traceback": exc.traceback}, + ) from exc + except BubError as exc: + raise BubError(exc.kind, exc.message, details={"output": "".join(output), **(exc.details or {})}) from exc + finally: + _ACTIVE_RUN.reset(token) + await republic_session.aclose() + if media: + return content_result(["".join(output), *media] if output else media) + return "".join(output) + + +_RUN_CODE_DESCRIPTION = """Run Python code in the environment and return everything it prints. + +Tools are async functions available as `tools.(...)`: await them with keyword arguments +(top-level `await` is allowed). See the tool stub referenced in the system prompt for their +signatures and result types. The `republic` module provides media loaders and models. Media passed +to `tools.attach_image/audio/video`, as an object, path or URL, is attached to the result. The code is stopped after timeout_seconds.""" + + +def write_stub(content: str, directory: Path) -> Path: + """Write a stub into ``directory``; the file name depends only on the content, so unchanged stubs keep their path.""" + path = directory / f"tools-{hashlib.sha256(content.encode()).hexdigest()[:12]}.pyi" + if not path.exists(): + path.parent.mkdir(parents=True, exist_ok=True) + temp_path = path.with_name(f".{path.name}.{uuid.uuid4().hex}.tmp") + temp_path.write_text(content, encoding="utf-8") + temp_path.replace(path) + return path + + +def render_code_mode_prompt(stub: Path | str, *, tool_name: str = RUN_CODE_TOOL_NAME) -> str: + """Render the code mode prompt for a stub file path, or for the stub text itself.""" + if isinstance(stub, Path): + where = f"in the stub file: {stub}\nRead the stub before calling a tool you have not used yet. " + else: + where = f"in this stub:\n```python\n{stub.rstrip()}\n```\n" + return ( + "\n" + f"More tools are available as async Python functions `tools.(...)` inside `{tool_name}`. " + f"Their signatures, result types and documentation are {where}" + "Always `await` tool calls (top-level `await` " + "is allowed) and pass keyword arguments; they return structured values and raise on failure. " + f"`{tool_name}` returns only what the code prints, so print the results you need, and combine " + "several tool calls in one run when possible. To look at an image, audio or video, pass its path, URL " + "or `republic` media object to `tools.attach_image(image=...)`, `tools.attach_audio(audio=...)` or " + f"`tools.attach_video(video=...)`: it is attached to the `{tool_name}` result.\n" + "The code can also use the `republic` module (see the end of the stub) to call models with " + "`republic.get_model()`, " + "`republic.get_embedding_model(spec)` and `republic.get_decision_model(spec)`; `get_model()` without a " + "spec is the default model of this session.\n" + "" + ) + + +def _code_tool_caller( + tools_by_name: Mapping[str, Tool], context: ToolContext, republic_session: RepublicSession +) -> CallTool: + """Serve the code's ``tools.`` and ``republic`` calls.""" + # Code consumes structured results, so this executor does not render them to text. + executor = ToolExecutor(context.events, render=False) + + async def call_tool(name: str, arguments: dict[str, Any]) -> Any: + arguments = decode_value(arguments) + if name.startswith(REPUBLIC_CALL_PREFIX): + return encode_value(await republic_session.call(name, arguments)) + tool_item = tools_by_name.get(name) + if tool_item is None: + raise BubError(ErrorKind.INVALID_INPUT, f"unknown tool: tools.{name}") + return encode_value(await executor.execute(tool_item, arguments, context=context)) + + return call_tool + + +_REMOTE_MEDIA_PREFIXES = ("http://", "https://", "gs://", "data:") +_MEDIA_CLASSES: dict[str, type[republic.Image | republic.Audio | republic.Video]] = { + "image": republic.Image, + "audio": republic.Audio, + "video": republic.Video, +} + + +async def _load_media( + kind: str, source: str, environment: Environment +) -> republic.Image | republic.Audio | republic.Video: + """Load media from a URL, a data URL, or a file path in the environment.""" + if source.startswith(_REMOTE_MEDIA_PREFIXES): + media: republic.Image | republic.Audio | republic.Video = getattr(republic, kind)(source) + return media + media_type, _ = mimetypes.guess_type(source) + if media_type is None: + raise ValueError(f"cannot guess the media type of {source!r}") + data = await environment.read_bytes(environment.resolve_path(source)) + return _MEDIA_CLASSES[kind](media_type, data=data) + + +async def _attach(kind: str, item: republic.Image | republic.Audio | republic.Video | str) -> str: + run = _ACTIVE_RUN.get() + if run is None: + raise BubError(ErrorKind.INVALID_INPUT, f"tools.attach_{kind} can only be called from {RUN_CODE_TOOL_NAME}.") + if isinstance(item, str): + try: + item = await _load_media(kind, item, run.environment) + except (OSError, ValueError) as exc: + raise BubError(ErrorKind.INVALID_INPUT, f"Cannot load the {kind}: {exc}") from exc + run.media.append(item) + return f"The {kind} is attached to the {RUN_CODE_TOOL_NAME} result." + + +async def _attach_image(image: republic.Image | str, *, context: ToolContext) -> str: + """Attach an image to the result of this `run_code` call, so you can see it. + + `image` is a `republic.Image`, a file path in the environment, an http(s) URL, or a data URL. + """ + return await _attach("image", image) + + +async def _attach_audio(audio: republic.Audio | str, *, context: ToolContext) -> str: + """Attach audio to the result of this `run_code` call, so you can hear it. + + `audio` is a `republic.Audio`, a file path in the environment, an http(s) URL, or a data URL. + """ + return await _attach("audio", audio) + + +async def _attach_video(video: republic.Video | str, *, context: ToolContext) -> str: + """Attach a video to the result of this `run_code` call, so you can watch it. + + `video` is a `republic.Video`, a file path in the environment, an http(s) URL, or a data URL. + """ + return await _attach("video", video) + + +attach_image = Tool.from_callable(_attach_image, name="attach_image", context=True, exposure="code") +attach_audio = Tool.from_callable(_attach_audio, name="attach_audio", context=True, exposure="code") +attach_video = Tool.from_callable(_attach_video, name="attach_video", context=True, exposure="code") +MEDIA_TOOLS = (attach_image, attach_audio, attach_video) +"""Code-only tools attaching media to the ``run_code`` result; they work only inside ``run_code``.""" diff --git a/src/bub/builtin/codemode/republic_bridge.py b/src/bub/core/codemode/republic_bridge.py similarity index 87% rename from src/bub/builtin/codemode/republic_bridge.py rename to src/bub/core/codemode/republic_bridge.py index 902d3ceb..833744d6 100644 --- a/src/bub/builtin/codemode/republic_bridge.py +++ b/src/bub/core/codemode/republic_bridge.py @@ -1,7 +1,8 @@ -"""Serve the ``republic`` module of ``run_code``: Republic values over JSON and model calls made in Bub. +"""Serve the ``republic`` module of ``run_code``: Republic values over JSON and model calls made by the host. Model-written code never sees provider credentials: its models are proxies whose calls arrive -here as tool calls named ``republic.`` (see ``code_runner_child``). +here as tool calls named ``republic.`` (see ``runner_child``), and a +:data:`ModelResolver` turns their specs into Republic specs and client options. """ from __future__ import annotations @@ -35,6 +36,23 @@ REPUBLIC_CALL_PREFIX = "republic." REPUBLIC_KEY = "$republic" +ModelResolver = Callable[[str | None, ToolContext], tuple[str, Mapping[str, Any]]] +"""Resolve the spec the code passed to ``republic.get_*model`` (None for ``get_model()``) to a Republic +``provider:model`` spec and the client options to open it with. Raise ``ValueError`` to reject the spec.""" + + +def default_model_resolver(default: str | None = None) -> ModelResolver: + """Pass specs through, with ``default`` for ``get_model()``; clients read credentials from the environment.""" + + def resolve(spec: str | None, context: ToolContext) -> tuple[str, Mapping[str, Any]]: + spec = spec or default + if not spec: + raise ValueError("No default model: pass a 'provider:model' spec.") + return spec, {} + + return resolve + + _MEDIA_CLASSES: dict[str, type[republic.Image | republic.Audio | republic.Video]] = { "image": republic.Image, "audio": republic.Audio, @@ -98,21 +116,6 @@ def _encode_message(message: republic.Message) -> dict[str, Any]: return {REPUBLIC_KEY: "message", "role": message.role, "content": content} -def _model_spec(context: ToolContext, spec: str | None) -> tuple[str, dict[str, Any]]: - """Resolve a spec, or the session's model, to a Republic spec and client options from Bub's settings.""" - agent = context.state.get("_runtime_agent") - settings = getattr(agent, "settings", None) - if settings is None: - raise BubError(ErrorKind.INVALID_INPUT, "Republic models need a running Bub agent.") - spec = spec or context.state.get("model") or settings.model - try: - candidate = settings.model_candidates(spec)[0] - except ValueError as exc: - raise BubError(ErrorKind.INVALID_INPUT, str(exc)) from exc - client_kwargs = settings.model_client_kwargs(candidate.provider_name or candidate.provider) - return f"{candidate.provider}:{candidate.model_id}", client_kwargs - - def _chat_messages(messages: list[Any], options: dict[str, Any]) -> list[republic.Message]: if "tools" in options or "tool_choice" in options: raise BubError(ErrorKind.INVALID_INPUT, "republic models in run_code do not support tools.") @@ -196,8 +199,9 @@ async def aclose(self) -> None: class RepublicSession: """Serve the ``republic.`` calls of one ``run_code``; ``aclose`` closes streams left open.""" - def __init__(self, context: ToolContext) -> None: + def __init__(self, context: ToolContext, resolve_model: ModelResolver) -> None: self.context = context + self._resolve_model = resolve_model self._streams: dict[str, _StreamPump] = {} self._operations: dict[str, Callable[..., Awaitable[Any]]] = { "chat": self._chat, @@ -222,20 +226,26 @@ async def call(self, name: str, arguments: dict[str, Any]) -> Any: except Exception as exc: raise BubError(ErrorKind.TOOL, f"{name} failed: {type(exc).__name__}: {exc}") from exc + def _model_spec(self, spec: str | None) -> tuple[str, Mapping[str, Any]]: + try: + return self._resolve_model(spec, self.context) + except ValueError as exc: + raise BubError(ErrorKind.INVALID_INPUT, str(exc)) from exc + async def aclose(self) -> None: streams, self._streams = list(self._streams.values()), {} await asyncio.gather(*(stream.aclose() for stream in streams), return_exceptions=True) async def _chat(self, spec: str | None, messages: list[Any], options: dict[str, Any]) -> Any: messages = _chat_messages(messages, options) - model_spec, client_kwargs = _model_spec(self.context, spec) + model_spec, client_kwargs = self._model_spec(spec) async with republic.get_model(model_spec, **client_kwargs) as model: response = await model.chat(messages, **options) return _encode_response(response) async def _stream_open(self, spec: str | None, messages: list[Any], options: dict[str, Any]) -> Any: messages = _chat_messages(messages, options) - model_spec, client_kwargs = _model_spec(self.context, spec) + model_spec, client_kwargs = self._model_spec(spec) async with AsyncExitStack() as stack: model = await stack.enter_async_context(republic.get_model(model_spec, **client_kwargs)) stream = await stack.enter_async_context(model.stream(messages, **options)) @@ -257,7 +267,7 @@ async def _stream_close(self, stream: str) -> None: await pump.aclose() async def _embed(self, spec: str, texts: list[str], dimensions: int | None = None) -> Any: - model_spec, client_kwargs = _model_spec(self.context, spec) + model_spec, client_kwargs = self._model_spec(spec) async with republic.get_embedding_model(model_spec, **client_kwargs) as model: response = await model.embed_many(texts, dimensions=dimensions) return {"vectors": response.vectors, "model": response.model, "token_usage": asdict(response.token_usage)} @@ -265,7 +275,7 @@ async def _embed(self, spec: str, texts: list[str], dimensions: int | None = Non async def _decide(self, spec: str, state: Any, questions: dict[str, Question]) -> Any: if not all(isinstance(question, Noul | Choice | Score) for question in questions.values()): raise BubError(ErrorKind.INVALID_INPUT, "decide() takes republic.decisions questions.") - model_spec, client_kwargs = _model_spec(self.context, spec) + model_spec, client_kwargs = self._model_spec(spec) async with republic.get_decision_model(model_spec, **client_kwargs) as model: response = await model.decide(state, questions=questions) return { diff --git a/src/bub/builtin/codemode/code_runner.py b/src/bub/core/codemode/runner.py similarity index 96% rename from src/bub/builtin/codemode/code_runner.py rename to src/bub/core/codemode/runner.py index faffeb73..f9c86156 100644 --- a/src/bub/builtin/codemode/code_runner.py +++ b/src/bub/core/codemode/runner.py @@ -1,6 +1,6 @@ """Run ``Environment.run_code`` code in a Python process started with ``Environment.spawn``. -The process runs ``code_runner_child`` and talks to Bub over stdin/stdout with +The process runs ``runner_child`` and talks to Bub over stdin/stdout with JSON lines; see that module for the protocol. """ @@ -17,7 +17,7 @@ from pydantic_core import to_jsonable_python -from bub.builtin.codemode import code_runner_child +from bub.core.codemode import runner_child from bub.environment import CodeFailed from bub.errors import BubError, ErrorKind @@ -25,7 +25,7 @@ from bub.environment import CallTool, Environment, Process _STOP_TIMEOUT_SECONDS = 3.0 -_CHILD_SOURCE = Path(code_runner_child.__file__).read_text(encoding="utf-8") +_CHILD_SOURCE = Path(runner_child.__file__).read_text(encoding="utf-8") class _Session: diff --git a/src/bub/builtin/codemode/code_runner_child.py b/src/bub/core/codemode/runner_child.py similarity index 100% rename from src/bub/builtin/codemode/code_runner_child.py rename to src/bub/core/codemode/runner_child.py diff --git a/src/bub/core/codemode/stub.py b/src/bub/core/codemode/stub.py new file mode 100644 index 00000000..f57726bc --- /dev/null +++ b/src/bub/core/codemode/stub.py @@ -0,0 +1,277 @@ +"""Render tools as a Python stub that tells model-written code their signatures and result types.""" + +from __future__ import annotations + +import keyword +import re +from collections.abc import Iterable +from typing import Any + +from bub.core.tools import Tool, model_tools + +_STUB_HEADER = '''"""Tools available inside `run_code` as `tools.(...)`. + +Every tool is an async function: `await` it (top-level `await` is allowed) and pass keyword +arguments only. Each call returns the structured value described by its return type and +raises an exception when the tool fails. Use `asyncio.gather` to run independent calls +concurrently. + +The `republic` module is also available (preset as a global and importable); it is documented +at the end of this file. +""" + +from typing import Any, Literal, NotRequired, TypedDict + +import republic''' + +_REPUBLIC_STUB = """# --------------------------------------------------------------------------- +# `republic`: media and models inside `run_code`. Models are served by the host with +# its credentials and settings; every model method is async. Only the members below exist. +# +# class republic.Image / republic.Audio / republic.Video: +# __init__(media_type: str, data: bytes | None = None, url: str | None = None) +# media_type: str; data: bytes | None; url: str | None +# base64_data: str; data_url: str +# `tools.attach_image/audio/video` also take a file path in the environment, an http(s)/gs URL +# or a data URL instead, and load it for you. +# +# republic.system(text: str) -> Message +# republic.user(*content: str | Image | Audio | Video) -> Message +# republic.assistant(*content: str | Image | Audio | Video) -> Message +# Message has .role, .content (list of text and media) and .text. +# +# republic.get_model(spec: str | None = None) -> ChatModel +# A "provider:model" spec; None is the model running this session. +# await model.chat(prompt: str | Message | list[str | Message], **options) -> response +# options: max_tokens, temperature, top_p, top_k, stop, seed, reasoning_effort, extra_body... +# (no tools). A plain string is a user message. +# response.text, .reasoning, .refusal, .finish_reason, .model, +# .token_usage.input_tokens / .output_tokens, .message (append it to continue the conversation) +# model.stream(prompt, **options) -> Stream # same arguments as chat(); not awaited +# async with model.stream(...) as stream: +# async for event in stream: ... +# Events from republic.events: TextDelta(.chunk), ReasoningDelta(.chunk), RefusalDelta(.chunk), +# UsageDelta(.usage), ImageReady(.image), CitationAdded(.citation), Completed(.response, last). +# After iterating: stream.response (as returned by chat()), stream.text, .reasoning, .token_usage. +# republic.get_embedding_model(spec: str) -> EmbeddingModel +# await model.embed(text: str, *, dimensions: int | None = None) -> response # response.vector +# await model.embed_many(texts: list[str], *, dimensions: int | None = None) -> response # response.vectors +# republic.get_decision_model(spec: str) -> DecisionModel +# await model.decide(state, *, questions: dict[str, Question]) -> response +# Questions from republic.decisions: Noul(instructions, criteria=None) (yes/no), +# Choice(instructions, criteria) (pick one), Score(instructions, criteria) (ordered levels). +# response.answers[id] or response.: .noul | .choice/.probabilities/.confidence | .score/... +# +# Example: +# model = republic.get_model() +# chart = republic.Image("image/png", data=open("chart.png", "rb").read()) +# reply = await model.chat([republic.user("Describe this chart.", chart)]) +# print(reply.text) +# async with model.stream("Write a haiku.") as stream: +# async for event in stream: +# if isinstance(event, republic.events.TextDelta): +# print(event.chunk, end="")""" +_JSON_TYPES = {"string": "str", "integer": "int", "number": "float", "boolean": "bool", "null": "None"} + + +def _identifier(name: str) -> str: + ident = re.sub(r"\W", "_", name) + if not ident or ident[0].isdigit() or keyword.iskeyword(ident): + ident = f"_{ident}" + return ident + + +def _is_identifier(name: str) -> bool: + return name.isidentifier() and not keyword.iskeyword(name) + + +def _pascal(name: str) -> str: + return "".join(part[:1].upper() + part[1:] for part in re.split(r"[^0-9A-Za-z]+", name)) + + +_MEDIA_FIELDS = {"media_type", "data", "url"} + + +def _media_class(schema: dict[str, Any]) -> str | None: + """Recognize the JSON schema of ``republic.Image``, ``Audio`` and ``Video`` parameters.""" + title = schema.get("title") + properties = schema.get("properties") + if title in ("Image", "Audio", "Video") and isinstance(properties, dict) and set(properties) == _MEDIA_FIELDS: + return str(title) + return None + + +def _docstring(text: str, indent: str = " ") -> list[str]: + text = text.strip().replace("\\", "\\\\").replace('"""', '\\"\\"\\"') + if not text: + return [] + lines = text.splitlines() + if len(lines) == 1: + return [f'{indent}"""{lines[0]}"""'] + return [f'{indent}"""{lines[0]}', *(f"{indent}{line}".rstrip() for line in lines[1:]), f'{indent}"""'] + + +class _StubBuilder: + """Render tool JSON schemas as Python stub declarations.""" + + def __init__(self) -> None: + self.classes: dict[str, str] = {} + self._refs: dict[tuple[int, str], str] = {} + + def function(self, tool_item: Tool) -> str: + name = _identifier(tool_item.name) + schema = tool_item.parameters or {} + defs = self._defs(schema) + properties = schema.get("properties") or {} + required = set(schema.get("required") or []) + params: list[str] = [] + arg_docs: list[str] = [] + for key, prop in properties.items(): + hint = self.type_hint(prop, defs, f"{_pascal(name)}{_pascal(key)}") + if key in required: + params.append(f"{key}: {hint}") + elif isinstance(prop, dict) and "default" in prop: + params.append(f"{key}: {hint} = {prop['default']!r}") + else: + params.append(f"{key}: {hint} = ...") + description = prop.get("description") if isinstance(prop, dict) else None + if description: + arg_docs.append(f" {key}: {' '.join(description.split())}") + if not all(_is_identifier(key) for key in properties): + signature = "**kwargs: Any" + elif params: + signature = f"*, {', '.join(params)}" + else: + signature = "" + + output = tool_item.output_schema + returns = "Any" if output is None else self.type_hint(output, self._defs(output), f"{_pascal(name)}Result") + doc = tool_item.description.strip() + if arg_docs: + doc = f"{doc}\n\nArgs:\n" + "\n".join(arg_docs) if doc else "Args:\n" + "\n".join(arg_docs) + body = _docstring(doc) or [" ..."] + return "\n".join([f"async def {name}({signature}) -> {returns}:", *body]) + + def type_hint( # noqa: C901 + self, schema: Any, defs: dict[str, Any], name_hint: str, ref_key: tuple[int, str] | None = None + ) -> str: + if not isinstance(schema, dict) or not schema: + return "Any" + if (media := _media_class(schema)) is not None: + return f"republic.{media}" + if local_defs := self._defs(schema): + # Schemas of single parameters carry their own definitions. + defs = {**defs, **local_defs} + if isinstance(ref := schema.get("$ref"), str): + return self._ref(ref, defs) + for key in ("anyOf", "oneOf"): + if isinstance(options := schema.get(key), list): + return self._union(self.type_hint(option, defs, name_hint) for option in options) + if isinstance(all_of := schema.get("allOf"), list) and len(all_of) == 1: + return self.type_hint(all_of[0], defs, name_hint, ref_key) + if "const" in schema: + return f"Literal[{schema['const']!r}]" + if isinstance(enum := schema.get("enum"), list) and enum: + return f"Literal[{', '.join(repr(value) for value in enum)}]" + + schema_type = schema.get("type") + if isinstance(schema_type, list): + return self._union(self.type_hint({**schema, "type": item}, defs, name_hint) for item in schema_type) + if schema_type == "array": + if isinstance(prefix := schema.get("prefixItems"), list): + return f"tuple[{', '.join(self.type_hint(item, defs, name_hint) for item in prefix)}]" + return f"list[{self.type_hint(schema.get('items'), defs, f'{name_hint}Item')}]" + if schema_type == "object": + if isinstance(properties := schema.get("properties"), dict) and properties: + return self._typed_dict(schema, properties, defs, schema.get("title") or name_hint, ref_key) + extra = schema.get("additionalProperties") + if isinstance(extra, dict) and extra: + return f"dict[str, {self.type_hint(extra, defs, f'{name_hint}Value')}]" + return "dict[str, Any]" + return _JSON_TYPES.get(schema_type, "Any") if isinstance(schema_type, str) else "Any" + + @staticmethod + def _defs(schema: dict[str, Any]) -> dict[str, Any]: + defs = schema.get("$defs", schema.get("definitions")) + return defs if isinstance(defs, dict) else {} + + @staticmethod + def _union(hints: Iterable[str]) -> str: + unique = list(dict.fromkeys(hints)) + return "Any" if "Any" in unique else " | ".join(unique) + + def _ref(self, ref: str, defs: dict[str, Any]) -> str: + name = ref.rsplit("/", 1)[-1] + key = (id(defs), name) + if key in self._refs: + return self._refs[key] + target = defs.get(name) + if not isinstance(target, dict): + return "Any" + hint = self.type_hint(target, defs, name, ref_key=key) + self._refs[key] = hint + return hint + + def _typed_dict( + self, + schema: dict[str, Any], + properties: dict[str, Any], + defs: dict[str, Any], + title: str, + ref_key: tuple[int, str] | None, + ) -> str: + base = _identifier(_pascal(title) or "Result") + class_name = base + index = 2 + while class_name in self.classes: + class_name = f"{base}{index}" + index += 1 + # Reserve the name first so recursive references resolve to this class. + self.classes[class_name] = "" + if ref_key is not None: + self._refs[ref_key] = class_name + + required = set(schema.get("required") or []) + fields: list[tuple[str, str, str | None]] = [] + for key, prop in properties.items(): + hint = self.type_hint(prop, defs, f"{class_name}{_pascal(key)}") + if key not in required: + hint = f"NotRequired[{hint}]" + description = prop.get("description") if isinstance(prop, dict) else None + fields.append((key, hint, description)) + source = self._class_source(class_name, fields, schema.get("description")) + + # Reuse an identical class that was declared under the same title by another tool. + for existing, existing_source in self.classes.items(): + if ( + existing != class_name + and re.fullmatch(rf"{re.escape(base)}\d*", existing) + and existing_source == source.replace(class_name, existing) + ): + del self.classes[class_name] + if ref_key is not None: + self._refs[ref_key] = existing + return existing + self.classes[class_name] = source + return class_name + + @staticmethod + def _class_source(name: str, fields: list[tuple[str, str, str | None]], description: object) -> str: + if not all(_is_identifier(key) for key, _, _ in fields): + items = ", ".join(f"{key!r}: {hint}" for key, hint, _ in fields) + return f"{name} = TypedDict({name!r}, {{{items}}})" + lines = [f"class {name}(TypedDict):"] + if isinstance(description, str): + lines.extend(_docstring(description)) + for key, hint, field_description in fields: + if field_description: + lines.append(f" # {' '.join(field_description.split())}") + lines.append(f" {key}: {hint}") + return "\n".join(lines) + + +def render_tool_stub(tools: Iterable[Tool]) -> str: + """Render a Python stub declaring one function per model-facing tool.""" + builder = _StubBuilder() + functions = [builder.function(tool_item) for tool_item in model_tools(tools)] + return "\n\n\n".join([_STUB_HEADER, *builder.classes.values(), *functions, _REPUBLIC_STUB]) + "\n" diff --git a/src/bub/environment.py b/src/bub/environment.py index 6dab1024..242d5b3e 100644 --- a/src/bub/environment.py +++ b/src/bub/environment.py @@ -5,15 +5,20 @@ ``bash``, ``fs.*`` and ``run_code`` run on top of it. Tool semantics (background shells, timeouts, rendering, hooks) stay on the host; only execution moves into the environment. Plugins provide one per session through the ``provide_environment`` hook; Bub's -builtin hooks provide ``bub.builtin.environment.LocalEnvironment``, which runs -everything on the host. +builtin hooks provide :class:`LocalEnvironment`, which runs everything on the host. """ from __future__ import annotations import abc import asyncio +import contextlib +import os +import shutil +import signal +import sys from collections.abc import Awaitable, Callable, Mapping, Sequence +from pathlib import Path from typing import Any ENVIRONMENT_STATE_KEY = "_runtime_environment" @@ -140,16 +145,132 @@ async def run_code( The code may use top-level ``await``. For each name in ``tools``, ``tools.(**kwargs)`` must be an async function that returns ``await call_tool(name, kwargs)``; errors from ``call_tool`` surface in the code as exceptions. The code may also call ``call_tool`` with names - prefixed ``republic.``; ``bub.builtin.codemode.code_runner_child`` shows the ``republic`` module + prefixed ``republic.``; ``bub.core.codemode.runner_child`` shows the ``republic`` module that makes those calls. Pass everything the code writes to stdout to ``write`` as it happens, so output survives a timeout. Raise :class:`CodeFailed` when the code raises; any other exception reports that the runtime itself failed. When cancelled (for example on timeout), stop the code before returning. How the code runs is up to the environment: a Python process, an embedded interpreter or a - remote code interpreter. ``bub.builtin.codemode.code_runner.run_code_in_subprocess`` implements it + remote code interpreter. ``bub.core.codemode.runner.run_code_in_subprocess`` implements it on top of :meth:`spawn` for environments that have a Python interpreter. """ async def close(self) -> None: # noqa: B027 """Release the environment's resources. Called when the framework stops.""" + + +class LocalProcess(Process): + """A host process in its own session, so signals reach the whole process group.""" + + def __init__(self, process: asyncio.subprocess.Process) -> None: + self.process = process + + @property + def pid(self) -> int: + return self.process.pid + + @property + def stdout(self) -> asyncio.StreamReader: + assert self.process.stdout is not None # noqa: S101 + return self.process.stdout + + @property + def stderr(self) -> asyncio.StreamReader: + assert self.process.stderr is not None # noqa: S101 + return self.process.stderr + + @property + def returncode(self) -> int | None: + return self.process.returncode + + async def wait(self) -> int: + return await self.process.wait() + + async def write_stdin(self, data: bytes) -> None: + assert self.process.stdin is not None # noqa: S101 + self.process.stdin.write(data) + await self.process.stdin.drain() + + def close_stdin(self) -> None: + if self.process.stdin is not None: + self.process.stdin.close() + + def signal(self, *, kill: bool) -> None: + with contextlib.suppress(ProcessLookupError): + if os.name != "nt": + os.killpg(self.process.pid, signal.SIGKILL if kill else signal.SIGTERM) + elif self.returncode is None: + if kill: + self.process.kill() + else: + self.process.terminate() + + def is_running(self) -> bool: + if os.name == "nt": + return self.returncode is None + try: + os.killpg(self.process.pid, 0) + except ProcessLookupError: + return False + except PermissionError: + # EPERM does not establish that the group is gone (in particular + # while its leader is exiting on macOS). + return True + return True + + +class LocalEnvironment(Environment): + """Run processes and access files directly on the host.""" + + SHELL = shutil.which("bash") or shutil.which("sh") if os.name != "nt" else None + + def __init__(self, workspace: str | Path | None = None) -> None: + self.workspace = str(workspace) if workspace is not None else None + self.python = sys.executable + + async def spawn( + self, command: str | Sequence[str], *, cwd: str | None = None, env: Mapping[str, str] | None = None + ) -> LocalProcess: + options: dict[str, Any] = { + "cwd": cwd or self.workspace, + "env": {**os.environ, **env} if env else None, + "stdin": asyncio.subprocess.PIPE, + "stdout": asyncio.subprocess.PIPE, + "stderr": asyncio.subprocess.PIPE, + "start_new_session": os.name != "nt", + } + if isinstance(command, str): + process = await asyncio.create_subprocess_shell(command, executable=self.SHELL, **options) + else: + process = await asyncio.create_subprocess_exec(*command, **options) + return LocalProcess(process) + + async def read_text(self, path: str) -> str: + return await asyncio.to_thread(Path(path).read_text, encoding="utf-8") + + async def read_bytes(self, path: str) -> bytes: + return await asyncio.to_thread(Path(path).read_bytes) + + async def write_text(self, path: str, content: str) -> None: + def write() -> None: + target = Path(path) + target.parent.mkdir(parents=True, exist_ok=True) + target.write_text(content, encoding="utf-8") + + await asyncio.to_thread(write) + + async def run_code( + self, code: str, *, tools: Sequence[str], call_tool: CallTool, write: Callable[[str], None] + ) -> None: + from bub.core.codemode.runner import run_code_in_subprocess + + await run_code_in_subprocess(self, code, tools=tools, call_tool=call_tool, write=write) + + def resolve_path(self, path: str) -> str: + expanded = Path(path).expanduser() + if expanded.is_absolute(): + return str(expanded) + if self.workspace is None: + raise ValueError(f"relative path '{path}' is not allowed without a workspace") + return str((Path(self.workspace) / expanded).resolve()) diff --git a/src/bub/hooks/specs.py b/src/bub/hooks/specs.py index af4f3b58..7bcd7c06 100644 --- a/src/bub/hooks/specs.py +++ b/src/bub/hooks/specs.py @@ -139,7 +139,7 @@ def provide_environment(self, session_id: str, workspace: Path) -> Environment | Called once per session and cached until ``framework.running()`` exits, which closes it with ``Environment.close()``. ``workspace`` is the host - workspace. Bub's builtin hooks provide a host environment (``bub.builtin.environment.LocalEnvironment``); + workspace. Bub's builtin hooks provide a host environment (``bub.environment.LocalEnvironment``); a higher-priority plugin replaces it. """ diff --git a/tests/test_codemode.py b/tests/test_codemode.py index ff7f4e15..29cb4ec1 100644 --- a/tests/test_codemode.py +++ b/tests/test_codemode.py @@ -22,12 +22,11 @@ set_code_mode, write_tool_stub, ) -from bub.builtin.environment import LocalEnvironment, LocalProcess from bub.core.events import EventBus, ToolCallStarted from bub.core.store import AsyncTapeStoreAdapter, InMemoryTapeStore from bub.core.tape import Tape, TapeContext from bub.core.tools import REGISTRY, Tool, ToolContext, ToolExecutor, model_tools, tool -from bub.environment import ENVIRONMENT_STATE_KEY, CallTool, CodeFailed +from bub.environment import ENVIRONMENT_STATE_KEY, CallTool, CodeFailed, LocalEnvironment, LocalProcess from bub.errors import BubError, ErrorKind from bub.hooks.interception import ToolCallDecision @@ -588,8 +587,8 @@ async def test_run_code_republic_models_are_served_by_bub_settings( import republic from republic.decisions import Noul - from bub.builtin.codemode import republic_bridge from bub.builtin.settings import AgentSettings + from bub.core.codemode import republic_bridge _FakeChatModel.calls = [] for getter in ("get_model", "get_embedding_model", "get_decision_model"): @@ -673,8 +672,8 @@ def stream(self, messages: Any, **options: Any) -> _FakeStream: def _patch_streaming_model(monkeypatch: pytest.MonkeyPatch) -> SimpleNamespace: - from bub.builtin.codemode import republic_bridge from bub.builtin.settings import AgentSettings + from bub.core.codemode import republic_bridge _FakeStream.closed = [] monkeypatch.setattr(republic_bridge.republic, "get_model", lambda spec, **kwargs: _FakeStreamingModel()) diff --git a/tests/test_core_codemode.py b/tests/test_core_codemode.py new file mode 100644 index 00000000..4e84c507 --- /dev/null +++ b/tests/test_core_codemode.py @@ -0,0 +1,141 @@ +from __future__ import annotations + +import json +import subprocess +import sys +from pathlib import Path +from typing import Any + +import pytest +import republic + +from bub.core import Agent +from bub.core.codemode import CodeMode, default_model_resolver +from bub.core.events import ToolCallStarted +from bub.core.store import InMemoryTapeStore +from bub.core.tools import Tool, ToolContext, ToolExecutor +from bub.environment import LocalEnvironment +from tests.model_fakes import ProviderService, chat_events, sse, tool_events +from tests.test_core_boundary import APP_MODULES + + +def lookup(order_id: str) -> dict[str, Any]: + """Look up an order.""" + return {"id": order_id, "total": 12.5} + + +def _run_code_call(code: str) -> object: + arguments = json.dumps({"code": code}) + return sse(tool_events([{"index": 0, "id": "call-1", "function": {"name": "run_code", "arguments": arguments}}])) + + +def test_importing_core_codemode_loads_no_app_dependencies() -> None: + script = ( + "import sys, bub.core.codemode\n" + "loaded = {name.split('.')[0] for name in sys.modules}\n" + f"print(sorted(loaded & set({APP_MODULES!r})))\n" + "print(sorted(name for name in sys.modules if name.startswith(('bub.builtin', 'bub.framework'))))\n" + ) + result = subprocess.run([sys.executable, "-c", script], capture_output=True, text=True, check=True) + + assert result.stdout.splitlines() == ["[]", "[]"] + + +@pytest.mark.asyncio +async def test_core_agent_calls_tools_from_code(provider_service: ProviderService, tmp_path: Path) -> None: + provider_service.reply(_run_code_call("order = await tools.orders_get(order_id='A1')\nprint(order['total'])")) + provider_service.reply(sse(chat_events("It costs 12.5."))) + code_mode = CodeMode([Tool.from_callable(lookup, name="orders.get")], environment=LocalEnvironment(tmp_path)) + calls: list[tuple[str, bool]] = [] + + async with provider_service.client() as client: + model = republic.get_model("openai:test", http_client=client, api_format="chat", api_key="test") + agent = Agent( + store=InMemoryTapeStore(), tools=[code_mode.run_code], system_prompt=["Be brief.", code_mode.prompt] + ) + agent.events.on( + ToolCallStarted, + lambda event: calls.append((event.call.tool, event.call.context.code_mode)), # type: ignore[union-attr] + ) + async with agent.run(["How much is order A1?"], model=model, tape="session") as stream: + events = [event async for event in stream] + + assert events[-1].data == {"ok": True, "text": "It costs 12.5."} + first, second = provider_service.body(0), provider_service.body(1) + assert [tool["function"]["name"] for tool in first["tools"]] == ["run_code"] + system_prompt = first["messages"][0]["content"] + assert system_prompt.startswith("Be brief.\n\n") + assert "async def orders_get(*, order_id: str) -> dict[str, Any]:" in system_prompt + assert "async def attach_image(" in system_prompt + assert second["messages"][-1] == {"role": "tool", "tool_call_id": "call-1", "content": "12.5\n"} + assert calls == [("run_code", False), ("orders_get", True)] + + +def test_code_mode_prompt_points_at_a_stub_file_in_stub_dir(tmp_path: Path) -> None: + code_mode = CodeMode([Tool.from_callable(lookup)], stub_dir=tmp_path / "stubs", media_tools=False) + + prompt = code_mode.prompt(_prompt_context()) + + assert prompt is not None + stub_path = next((tmp_path / "stubs").iterdir()) + assert f"in the stub file: {stub_path}\n" in prompt + assert "async def lookup(*, order_id: str)" in stub_path.read_text(encoding="utf-8") + assert "attach_image" not in stub_path.read_text(encoding="utf-8").split("# `republic`")[0] + + +@pytest.mark.asyncio +async def test_code_mode_selects_tools_per_run_from_state(tmp_path: Path) -> None: + code_mode = CodeMode(lambda state: state.get("code_tools"), environment=LocalEnvironment(tmp_path)) + tape = Agent().tape("session") + + assert code_mode.prompt(_prompt_context()) is None + execution = await ToolExecutor().execute_many( + [(code_mode.run_code, {"code": "print(1)"})], context=ToolContext(tape=tape) + ) + assert execution.error is not None + assert execution.error.message == "Code mode is not enabled for this run." + + state = {"code_tools": [Tool.from_callable(lookup)]} + output = await code_mode.run_code.run( + code="print((await tools.lookup(order_id='B2'))['id'])", context=ToolContext(tape=tape, state=state) + ) + assert output == "B2\n" + assert [tool.name for tool in code_mode.code_tools(state) or []] == [ + "lookup", + "attach_image", + "attach_audio", + "attach_video", + ] + + +def test_default_model_resolver_needs_a_spec_or_default() -> None: + context = ToolContext(tape=Agent().tape("session")) + + assert default_model_resolver("openai:gpt-x")(None, context) == ("openai:gpt-x", {}) + assert default_model_resolver("openai:gpt-x")("anthropic:claude", context) == ("anthropic:claude", {}) + with pytest.raises(ValueError, match="No default model"): + default_model_resolver()(None, context) + + +@pytest.mark.asyncio +async def test_code_republic_models_use_the_resolver(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + def resolve(spec: str | None, context: ToolContext) -> tuple[str, dict[str, Any]]: + return f"resolved:{spec}", {"api_key": "secret"} + + def get_model(spec: str, **kwargs: Any) -> Any: + raise RuntimeError(f"{spec} {kwargs}") + + monkeypatch.setattr(republic, "get_model", get_model) + code_mode = CodeMode(environment=LocalEnvironment(tmp_path), models=resolve) + code = "try:\n await republic.get_model('x').chat('hi')\nexcept Exception as exc:\n print(exc)\n" + + output = await code_mode.run_code.run(code=code, context=ToolContext(tape=Agent().tape("session"))) + + assert output == "[tool] republic.chat failed: RuntimeError: resolved:x {'api_key': 'secret'}\n" + + +def _prompt_context() -> Any: + from bub.core import PromptContext + + agent = Agent() + return PromptContext(agent=agent, tape=agent.tape("session"), prompt=[], tools=[], model="openai:test") diff --git a/tests/test_environment.py b/tests/test_environment.py index 4d5b066d..77f0c5ae 100644 --- a/tests/test_environment.py +++ b/tests/test_environment.py @@ -6,12 +6,12 @@ import pytest import bub.builtin.tools as builtin_tools -from bub.builtin.environment import LocalEnvironment, LocalProcess, environment_from_state +from bub.builtin.environment import environment_from_state from bub.builtin.shell_manager import ShellManager from bub.core.store import AsyncTapeStoreAdapter, InMemoryTapeStore from bub.core.tape import Tape, TapeContext from bub.core.tools import ToolContext -from bub.environment import ENVIRONMENT_STATE_KEY, CallTool, Environment, Process +from bub.environment import ENVIRONMENT_STATE_KEY, CallTool, Environment, LocalEnvironment, LocalProcess, Process from bub.framework import BubFramework from bub.hooks import hookimpl diff --git a/website/src/content/docs/docs/build/tools.mdx b/website/src/content/docs/docs/build/tools.mdx index 506883b9..80a145a8 100644 --- a/website/src/content/docs/docs/build/tools.mdx +++ b/website/src/content/docs/docs/build/tools.mdx @@ -77,6 +77,30 @@ Nested calls run with `ToolContext.code_mode` set to `True` (hooks see it as `ca `run_code` hands the code to the session [environment](#run-tools-in-an-environment)'s `run_code`, together with a `call_tool` callback. Tool calls always run on the host through that callback, so hooks still see every call. `timeout_seconds` cancels the environment call. The builtin `LocalEnvironment` starts a fresh Python process for every call (`sys.executable`, working directory is the workspace) and forwards tool calls as JSON lines over its stdin and stdout, so arguments and results must be JSON-serializable. When the code finishes, fails or times out, it kills the process and everything it started. That process runs on the host with the same permissions as `bash`, so enable code mode only where `bash` would be acceptable. The return annotation of a tool function (or `Tool.output_schema`, for tools built by hand) determines the result type shown in the stub. +#### Code mode in your own agent + +Code mode lives in the kernel (`bub.core.codemode`), so an agent built on `bub.core.Agent` can use it without the `app` extra. A `CodeMode` holds the tools callable from code and provides the `run_code` tool and a system prompt section: + +```python +from bub.core import Agent +from bub.core.codemode import CodeMode +from bub.environment import LocalEnvironment + +code_mode = CodeMode( + [search, fetch], # tools callable as tools.(...) inside run_code + environment=LocalEnvironment("/work"), # where the code runs; defaults to the current directory + models="openai:gpt-5", # what republic.get_model() in the code resolves to +) +agent = Agent(model=model, tools=[code_mode.run_code], system_prompt=["You are ...", code_mode.prompt]) +``` + +- `tools` can also be a function of the turn state returning the tools for that run, or `None` to turn code mode off; `environment` can likewise be a function of the state. +- `models` is a default `provider:model` spec, or a resolver `(spec, context) -> (republic_spec, client_kwargs)` that keeps credentials on the host. Without it the code must pass explicit specs and clients read credentials from the environment. +- The stub goes into the system prompt itself; pass `stub_dir=` to write it to a file and reference the path instead, for agents that can read files. +- `media_tools=True` (default) adds `attach_image`, `attach_audio` and `attach_video`. + +Calls made from code go through the agent's `ToolCallStarted`/`ToolCallFinished` events with `context.code_mode` set, so handlers registered with `agent.events.on` see them. Bub's own agents use the same class (`bub.builtin.codemode.CODE_MODE`), wired to session settings, environments and model settings. + ### Load tools on demand Mark a tool `deferred=True` (on `@tool`, `Tool.from_callable`, or with `dataclasses.replace(tool, deferred=True)`) to keep its definition out of model requests until it is needed. Tools without the flag are loaded directly. @@ -94,10 +118,10 @@ Mark a tool `deferred=True` (on `@tool`, `Tool.from_callable`, or with `dataclas - `spawn(command, *, cwd=None, env=None)` returns a `Process`: a string runs through the shell, a sequence runs as an argument vector. The process exposes `stdout`/`stderr` streams, `write_stdin`, `wait`, and `signal(kill=...)`, which must reach the whole process tree. - `read_text(path)` and `write_text(path, content)` access files inside the environment. `read_bytes(path)` reads binary files; its default runs the environment's `python` to print the file, so override it when the environment has direct file access. - `workspace` is the working directory inside the environment; `resolve_path` resolves relative paths against it. -- `run_code(code, *, tools, call_tool, write)` runs Python code in which `await tools.(**kwargs)` calls `call_tool(name, kwargs)`. The `republic` module calls `call_tool` with names prefixed `republic.`, and Republic values travel as JSON objects with a `"$republic"` key; the subprocess runner provides it. It passes the code's stdout to `write` as it arrives, raises `CodeFailed` when the code raises, and stops the code when cancelled. How it runs is up to the environment. An environment that has a Python interpreter can reuse `bub.builtin.codemode.code_runner.run_code_in_subprocess`, which runs the code in a process started with `spawn`. +- `run_code(code, *, tools, call_tool, write)` runs Python code in which `await tools.(**kwargs)` calls `call_tool(name, kwargs)`. The `republic` module calls `call_tool` with names prefixed `republic.`, and Republic values travel as JSON objects with a `"$republic"` key; the subprocess runner provides it. It passes the code's stdout to `write` as it arrives, raises `CodeFailed` when the code raises, and stops the code when cancelled. How it runs is up to the environment. An environment that has a Python interpreter can reuse `bub.core.codemode.runner.run_code_in_subprocess`, which runs the code in a process started with `spawn`. - `close()` releases the environment. -Provide one per session with the `provide_environment(session_id, workspace)` hook. Bub calls it on the session's first turn, puts the result in `state["_runtime_environment"]`, and closes it when `framework.running()` exits. Bub's builtin hooks provide `LocalEnvironment` (in [`bub.builtin.environment`](https://github.com/bubbuild/bub/blob/main/src/bub/builtin/environment.py)), which runs everything on the host; a plugin's implementation takes precedence over it. Your own tools can use the same environment through `bub.builtin.environment.environment_from_state(context.state)`. +Provide one per session with the `provide_environment(session_id, workspace)` hook. Bub calls it on the session's first turn, puts the result in `state["_runtime_environment"]`, and closes it when `framework.running()` exits. Bub's builtin hooks provide `LocalEnvironment` (in [`bub.environment`](https://github.com/bubbuild/bub/blob/main/src/bub/environment.py)), which runs everything on the host; a plugin's implementation takes precedence over it. Your own tools can use the same environment through `bub.builtin.environment.environment_from_state(context.state)`. ## How tools become available diff --git a/website/src/content/docs/zh-cn/docs/build/tools.mdx b/website/src/content/docs/zh-cn/docs/build/tools.mdx index cf675b49..fa009623 100644 --- a/website/src/content/docs/zh-cn/docs/build/tools.mdx +++ b/website/src/content/docs/zh-cn/docs/build/tools.mdx @@ -78,6 +78,30 @@ Code mode 让模型在 Python 中调用工具，而不必一次只发起一个 `run_code` 会把代码和一个 `call_tool` 回调一起交给当前 session [执行环境](#在执行环境中运行工具)的 `run_code` 执行。工具调用始终经这个回调在宿主上执行，所以 hook 仍能看到每一次调用。超过 `timeout_seconds` 时会取消执行环境里的执行。builtin 的 `LocalEnvironment` 每次调用都会启动一个新的 Python 进程（使用 `sys.executable`，工作目录为 workspace），工具调用以 JSON 行的形式经进程的 stdin/stdout 转发，因此参数和结果都必须能 JSON 序列化。代码执行完、出错或超时后，它会杀掉该进程及其启动的所有子进程。这个进程直接跑在宿主上，权限与 `bash` 相同，因此只应在允许使用 `bash` 的场景下开启 code mode。stub 中的结果类型来自工具函数的返回注解（手动构造的工具则来自 `Tool.output_schema`）。 +#### 在自己的 agent 中使用 code mode + +Code mode 属于内核（`bub.core.codemode`），基于 `bub.core.Agent` 构建的 agent 不安装 `app` extra 也能使用。`CodeMode` 持有可在代码中调用的工具，并提供 `run_code` 工具和一个 system prompt 片段： + +```python +from bub.core import Agent +from bub.core.codemode import CodeMode +from bub.environment import LocalEnvironment + +code_mode = CodeMode( + [search, fetch], # 在 run_code 中以 tools.(...) 调用的工具 + environment=LocalEnvironment("/work"), # 代码运行的位置，默认是当前目录 + models="openai:gpt-5", # 代码里 republic.get_model() 对应的模型 +) +agent = Agent(model=model, tools=[code_mode.run_code], system_prompt=["You are ...", code_mode.prompt]) +``` + +- `tools` 也可以是以 turn state 为参数的函数，返回本次运行的工具，返回 `None` 表示关闭 code mode；`environment` 同样可以是 state 的函数。 +- `models` 是默认的 `provider:model` spec，或一个 resolver `(spec, context) -> (republic_spec, client_kwargs)`，凭据始终留在宿主侧。不传时代码必须显式给出 spec，客户端从环境变量读取凭据。 +- stub 默认直接放进 system prompt；传入 `stub_dir=` 则写入文件并在 prompt 中引用路径，适合能读文件的 agent。 +- `media_tools=True`（默认）会额外提供 `attach_image`、`attach_audio` 和 `attach_video`。 + +代码发起的调用同样经过 agent 的 `ToolCallStarted`/`ToolCallFinished` 事件，且 `context.code_mode` 为真，因此用 `agent.events.on` 注册的处理器都能看到。Bub 自己的 agent 使用的也是这个类（`bub.builtin.codemode.CODE_MODE`），只是接入了 session 设置、执行环境和模型配置。 + ### 按需加载工具 给工具标记 `deferred=True`（可用于 `@tool`、`Tool.from_callable`，或 `dataclasses.replace(tool, deferred=True)`），它的定义就不会进入模型请求，直到被加载。未标记的工具直接加载。 @@ -95,10 +119,10 @@ Code mode 让模型在 Python 中调用工具，而不必一次只发起一个 - `spawn(command, *, cwd=None, env=None)` 返回一个 `Process`：传字符串时经 shell 执行，传序列时作为参数列表直接执行。进程提供 `stdout`/`stderr` 流、`write_stdin`、`wait` 和 `signal(kill=...)`，其中 `signal` 必须作用到整个进程树。 - `read_text(path)` 和 `write_text(path, content)` 读写执行环境内的文件。`read_bytes(path)` 读取二进制文件；默认实现会用执行环境的 `python` 输出文件内容，能直接访问文件的执行环境应覆盖它。 - `workspace` 是执行环境内的工作目录，`resolve_path` 基于它解析相对路径。 -- `run_code(code, *, tools, call_tool, write)` 执行 Python 代码，代码里的 `await tools.(**kwargs)` 会调用 `call_tool(name, kwargs)`。`republic` 模块会以 `republic.` 前缀的名字调用 `call_tool`，Republic 值以带 `"$republic"` 键的 JSON 对象传输；子进程 runner 已经提供了这个模块。代码写到 stdout 的内容要实时交给 `write`；代码抛异常时抛出 `CodeFailed`；被取消时要停止代码。具体怎么执行由执行环境决定。有 Python 解释器的执行环境可以直接复用 `bub.builtin.codemode.code_runner.run_code_in_subprocess`，它会用 `spawn` 启动进程来执行代码。 +- `run_code(code, *, tools, call_tool, write)` 执行 Python 代码，代码里的 `await tools.(**kwargs)` 会调用 `call_tool(name, kwargs)`。`republic` 模块会以 `republic.` 前缀的名字调用 `call_tool`，Republic 值以带 `"$republic"` 键的 JSON 对象传输；子进程 runner 已经提供了这个模块。代码写到 stdout 的内容要实时交给 `write`；代码抛异常时抛出 `CodeFailed`；被取消时要停止代码。具体怎么执行由执行环境决定。有 Python 解释器的执行环境可以直接复用 `bub.core.codemode.runner.run_code_in_subprocess`，它会用 `spawn` 启动进程来执行代码。 - `close()` 释放执行环境。 -通过 `provide_environment(session_id, workspace)` hook 为每个 session 提供执行环境。Bub 在 session 的第一个 turn 调用它，把结果放进 `state["_runtime_environment"]`，并在 `framework.running()` 退出时关闭。Bub 的 builtin hook 提供 `LocalEnvironment`（见 [`bub.builtin.environment`](https://github.com/bubbuild/bub/blob/main/src/bub/builtin/environment.py)），直接在宿主上运行；插件提供的实现会优先于它。自定义工具可以用 `bub.builtin.environment.environment_from_state(context.state)` 使用同一个环境。 +通过 `provide_environment(session_id, workspace)` hook 为每个 session 提供执行环境。Bub 在 session 的第一个 turn 调用它，把结果放进 `state["_runtime_environment"]`，并在 `framework.running()` 退出时关闭。Bub 的 builtin hook 提供 `LocalEnvironment`（见 [`bub.environment`](https://github.com/bubbuild/bub/blob/main/src/bub/environment.py)），直接在宿主上运行；插件提供的实现会优先于它。自定义工具可以用 `bub.builtin.environment.environment_from_state(context.state)` 使用同一个环境。 ## 工具如何变得可用