adk-code-mode 0.1.0__py3-none-any.whl

adk_code_mode/__about__.py

@@ -0,0 +1,4 @@
+# SPDX-FileCopyrightText: 2025-present A2A Net <hello@a2anet.com>
+#
+# SPDX-License-Identifier: Apache-2.0
+__version__ = "0.1.0"  # x-release-please-version

adk_code_mode/__init__.py

@@ -0,0 +1,22 @@
+# SPDX-FileCopyrightText: 2025-present A2A Net <hello@a2anet.com>
+#
+# SPDX-License-Identifier: Apache-2.0
+from adk_code_mode.__about__ import __version__
+from adk_code_mode.callback import code_mode_before_model_callback
+from adk_code_mode.executor import (
+    CODE_MODE_SYSTEM_INSTRUCTION,
+    ArtifactsSavedCallback,
+    CodeModeCodeExecutor,
+)
+from adk_code_mode.runtime import DockerRuntime, SandboxHandle, SandboxRuntime
+
+__all__ = [
+    "ArtifactsSavedCallback",
+    "CODE_MODE_SYSTEM_INSTRUCTION",
+    "CodeModeCodeExecutor",
+    "DockerRuntime",
+    "SandboxHandle",
+    "SandboxRuntime",
+    "__version__",
+    "code_mode_before_model_callback",
+]

adk_code_mode/_artifact_tools.py

@@ -0,0 +1,152 @@
+# SPDX-FileCopyrightText: 2025-present A2A Net <hello@a2anet.com>
+#
+# SPDX-License-Identifier: Apache-2.0
+"""Built-in ``save_artifact`` / ``load_artifact`` / ``list_artifacts`` tools.
+
+These are regular ``FunctionTool`` instances injected at the front of
+``CodeModeCodeExecutor.tools`` (set ``include_artifact_tools=False`` to opt out).
+They appear in the rendered catalog as top-level tools; the model writes
+``from tools import save_artifact, load_artifact, list_artifacts``.
+
+The wire format is JSON. Binary content is base64 over the wire — the model
+encodes it before calling ``save_artifact`` and decodes it after
+``load_artifact``. Loaded artifacts return a flat ``{"kind", "data",
+"mime_type"}`` dict so the model dispatches on ``kind`` without importing
+host-side types.
+"""
+
+from __future__ import annotations
+
+import base64
+from typing import Any
+
+from google.adk.tools.function_tool import FunctionTool
+from google.adk.tools.tool_context import ToolContext
+from google.genai import types as genai_types
+
+_TEXT_MIME_PREFIXES = ("text/",)
+_TEXT_MIME_TYPES = {"application/json", "application/xml", "application/javascript"}
+
+
+def _is_text_mime(mime_type: str | None) -> bool:
+    if not mime_type:
+        return False
+    if mime_type.startswith(_TEXT_MIME_PREFIXES):
+        return True
+    if mime_type in _TEXT_MIME_TYPES:
+        return True
+    return mime_type.endswith("+json") or mime_type.endswith("+xml")
+
+
+async def save_artifact(
+    *,
+    filename: str,
+    content: str,
+    mime_type: str | None = None,
+    tool_context: ToolContext,
+) -> int:
+    """Save an artifact to the session. Returns the new version number.
+
+    For text or JSON, pass ``content`` as a string and set ``mime_type``:
+
+        save_artifact(
+            filename="report.json",
+            content=json.dumps({"status": "ok"}),
+            mime_type="application/json",
+        )
+
+    For binary, base64-encode the bytes and set the binary mime type:
+
+        save_artifact(
+            filename="image.png",
+            content=base64.b64encode(image_bytes).decode("ascii"),
+            mime_type="image/png",
+        )
+    """
+    if _is_text_mime(mime_type):
+        part = genai_types.Part(
+            inline_data=genai_types.Blob(
+                data=content.encode("utf-8"),
+                mime_type=mime_type or "text/plain",
+            )
+        )
+    else:
+        part = genai_types.Part(
+            inline_data=genai_types.Blob(
+                data=base64.b64decode(content),
+                mime_type=mime_type or "application/octet-stream",
+            )
+        )
+    return await tool_context.save_artifact(filename, part)
+
+
+async def load_artifact(
+    *,
+    filename: str,
+    version: int | None = None,
+    tool_context: ToolContext,
+) -> dict[str, Any] | None:
+    """Load an artifact by filename. Returns ``None`` if not found.
+
+    Returned dict shape:
+
+        {
+            "kind": "text" | "bytes",
+            "data": <str>,                # text: the string; bytes: base64-encoded
+            "mime_type": <str | None>,
+        }
+
+    Example:
+
+        result = load_artifact(filename="report.json")
+        if result is None:
+            return
+        if result["kind"] == "text" and result["mime_type"] == "application/json":
+            payload = json.loads(result["data"])
+        elif result["kind"] == "bytes":
+            blob = base64.b64decode(result["data"])
+    """
+    part = await tool_context.load_artifact(filename, version=version)
+    if part is None:
+        return None
+    inline = part.inline_data
+    if inline is not None:
+        mime_type = inline.mime_type
+        data = inline.data or b""
+        if isinstance(data, str):
+            data = data.encode("utf-8")
+        if _is_text_mime(mime_type):
+            return {
+                "kind": "text",
+                "data": data.decode("utf-8"),
+                "mime_type": mime_type,
+            }
+        return {
+            "kind": "bytes",
+            "data": base64.b64encode(data).decode("ascii"),
+            "mime_type": mime_type,
+        }
+    if part.text is not None:
+        return {"kind": "text", "data": part.text, "mime_type": None}
+    return {"kind": "text", "data": "", "mime_type": None}
+
+
+async def list_artifacts(*, tool_context: ToolContext) -> list[str]:
+    """Return the filenames of artifacts visible to this session.
+
+    Example:
+
+        for filename in list_artifacts():
+            print(filename)
+    """
+    return await tool_context.list_artifacts()
+
+
+ARTIFACT_TOOLS: tuple[FunctionTool, ...] = (
+    FunctionTool(save_artifact),
+    FunctionTool(load_artifact),
+    FunctionTool(list_artifacts),
+)
+
+
+__all__ = ["ARTIFACT_TOOLS", "list_artifacts", "load_artifact", "save_artifact"]

adk_code_mode/callback.py

@@ -0,0 +1,81 @@
+# SPDX-FileCopyrightText: 2025-present A2A Net <hello@a2anet.com>
+#
+# SPDX-License-Identifier: Apache-2.0
+"""``code_mode_before_model_callback``: inject the tool catalog into the prompt.
+
+Wire the result of :func:`code_mode_before_model_callback` as
+``before_model_callback`` on an ``LlmAgent``. On every model turn, the
+callback:
+
+1. Resolves the tools (including any ``BaseToolset`` instances) for the
+   current invocation, using the live ``CallbackContext``.
+2. Renders the catalog — signatures + docstrings for every tool —
+   falling back to a short overflow message when the rendered catalog
+   would exceed ``max_catalog_chars``.
+3. Appends ``\\n\\n<tools>\\n…\\n</tools>`` to
+   ``llm_request.config.system_instruction``.
+4. Caches the resolved tools on the executor so the follow-up
+   ``execute_code`` call reuses them instead of re-resolving toolsets.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Any, Awaitable, Callable
+
+from adk_code_mode.tools import namespacing, normaliser
+from adk_code_mode.tools.catalog import render_catalog, render_overflow_catalog
+
+if TYPE_CHECKING:
+    from adk_code_mode.executor import CodeModeCodeExecutor
+
+logger = logging.getLogger("adk_code_mode.callback")
+
+_TOOLS_OPEN = "<tools>"
+_TOOLS_CLOSE = "</tools>"
+
+
+def code_mode_before_model_callback(
+    executor: "CodeModeCodeExecutor",
+) -> Callable[..., Awaitable[None]]:
+    """Build a ``before_model_callback`` bound to ``executor``."""
+
+    async def _callback(
+        callback_context: Any,
+        llm_request: Any,
+        **_unused: Any,
+    ) -> None:
+        resolved = await normaliser.resolve(list(executor.tools), callback_context)
+
+        executor._record_resolved_tools(callback_context.invocation_id, resolved)
+
+        ns_tools = namespacing.build(resolved)
+        catalog = render_catalog(ns_tools)
+        if len(catalog) > executor.max_catalog_chars:
+            catalog = render_overflow_catalog()
+        block = f"{_TOOLS_OPEN}\n{catalog}\n{_TOOLS_CLOSE}"
+
+        _append_system_instruction(llm_request.config, block)
+
+    return _callback
+
+
+def _append_system_instruction(config: Any, block: str) -> None:
+    """Append ``block`` to ``config.system_instruction`` in place.
+
+    Handles the two shapes ADK exposes: ``None`` or ``str``, and a list of
+    items (typically strings or ``Part`` objects). The new block is always
+    appended as a final string; non-string items in a list are preserved
+    as-is.
+    """
+    existing = config.system_instruction
+    if existing is None:
+        config.system_instruction = block
+        return
+    if isinstance(existing, str):
+        config.system_instruction = f"{existing}\n\n{block}"
+        return
+    config.system_instruction = [*existing, block]
+
+
+__all__ = ["code_mode_before_model_callback"]