databricks-agentbricks 0.2.0__py3-none-any.whl

databricks_agentkit/runtime/store.py

@@ -0,0 +1,287 @@
+"""Runtime Store contracts and process-local implementation."""
+
+from __future__ import annotations
+
+import asyncio
+import copy
+import os
+from typing import Protocol
+
+from databricks_agentkit.runtime.types import (
+    Invocation,
+    InvocationConflictError,
+    InvocationEvent,
+    InvocationStatus,
+    JsonObject,
+    JsonValue,
+)
+
+DEFAULT_RUNTIME_STORE_SCHEMA = "databricks_agentkit_runtime"
+RUNTIME_STORE_LAKEBASE_BRANCH_ENV = "DATABRICKS_AGENTBRICKS_RUNTIME_STORE_LAKEBASE_BRANCH"
+RUNTIME_STORE_LAKEBASE_ENDPOINT_ENV = "DATABRICKS_AGENTBRICKS_RUNTIME_STORE_LAKEBASE_ENDPOINT"
+RUNTIME_STORE_DATABASE_ENV = "DATABRICKS_AGENTBRICKS_RUNTIME_STORE_DATABASE"
+RUNTIME_STORE_USERNAME_ENV = "DATABRICKS_AGENTBRICKS_RUNTIME_STORE_USERNAME"
+RUNTIME_STORE_SCHEMA_ENV = "DATABRICKS_AGENTBRICKS_RUNTIME_STORE_SCHEMA"
+RUNTIME_STORE_LOCAL_ENV = "DATABRICKS_AGENTBRICKS_RUNTIME_STORE_LOCAL"
+
+
+def _validate_invocation_id(invocation_id: str) -> None:
+    if not invocation_id:
+        raise ValueError("invocation_id must not be empty")
+
+
+class RuntimeStore(Protocol):
+    """Invocation state operations shared by local and durable Runtime Stores."""
+
+    async def initialize(self) -> None:
+        """Open the store and create any resources it owns."""
+        ...
+
+    async def close(self) -> None:
+        """Release the store's connections and other resources."""
+        ...
+
+    async def accept(self, invocation_id: str, request: JsonValue) -> Invocation:
+        """Create an invocation or return the identical request accepted for this ID.
+
+        Reusing an invocation ID with a different request raises :class:`InvocationConflictError`.
+        """
+        ...
+
+    async def get(self, invocation_id: str) -> Invocation | None:
+        """Return the latest invocation state, or ``None`` for an unknown ID."""
+        ...
+
+    async def claim(self, invocation_id: str) -> Invocation | None:
+        """Claim a newly queued invocation for its first attempt."""
+        ...
+
+    async def complete(
+        self,
+        invocation_id: str,
+        attempt: int,
+        response: JsonValue,
+    ) -> bool:
+        """Persist a result when this attempt still owns the active invocation."""
+        ...
+
+    async def fail(self, invocation_id: str, attempt: int) -> bool:
+        """Mark an invocation failed when this attempt still owns it."""
+        ...
+
+    async def append_event(
+        self,
+        invocation_id: str,
+        attempt: int,
+        event: JsonObject,
+    ) -> int | None:
+        """Append an event for an owned attempt and return its replay cursor."""
+        ...
+
+    async def events(
+        self,
+        invocation_id: str,
+        after_sequence: int | None = None,
+    ) -> list[InvocationEvent]:
+        """Return invocation events after the optional exclusive replay cursor."""
+        ...
+
+
+class InMemoryRuntimeStore(RuntimeStore):
+    """Process-local Runtime Store for development and tests."""
+
+    def __init__(self) -> None:
+        self.states: dict[str, Invocation] = {}
+        self.persisted_events: list[InvocationEvent] = []
+        self._lock = asyncio.Lock()
+
+    async def initialize(self) -> None:
+        pass
+
+    async def close(self) -> None:
+        pass
+
+    async def accept(self, invocation_id: str, request: JsonValue) -> Invocation:
+        _validate_invocation_id(invocation_id)
+        async with self._lock:
+            existing = self.states.get(invocation_id)
+            if existing is not None:
+                if existing.request != request:
+                    raise InvocationConflictError(invocation_id)
+                return copy.deepcopy(existing)
+            state = Invocation(
+                invocation_id=invocation_id,
+                status=InvocationStatus.QUEUED,
+                attempt=0,
+                request=copy.deepcopy(request),
+                response=None,
+            )
+            self.states[invocation_id] = state
+            return copy.deepcopy(state)
+
+    async def get(self, invocation_id: str) -> Invocation | None:
+        _validate_invocation_id(invocation_id)
+        async with self._lock:
+            state = self.states.get(invocation_id)
+            return copy.deepcopy(state) if state is not None else None
+
+    async def claim(self, invocation_id: str) -> Invocation | None:
+        _validate_invocation_id(invocation_id)
+        async with self._lock:
+            state = self.states.get(invocation_id)
+            if state is None or state.status != InvocationStatus.QUEUED:
+                return None
+            claimed = Invocation(
+                invocation_id=invocation_id,
+                status=InvocationStatus.ACTIVE,
+                attempt=state.attempt + 1,
+                request=copy.deepcopy(state.request),
+                response=None,
+            )
+            self.states[invocation_id] = claimed
+            self._append_event(invocation_id, claimed.attempt, {"type": "run.started"})
+            return copy.deepcopy(claimed)
+
+    async def complete(self, invocation_id: str, attempt: int, response: JsonValue) -> bool:
+        async with self._lock:
+            state = self.states.get(invocation_id)
+            if state is None or not self._owns_attempt(state, attempt):
+                return False
+            self.states[invocation_id] = Invocation(
+                invocation_id=state.invocation_id,
+                status=InvocationStatus.COMPLETED,
+                attempt=state.attempt,
+                request=state.request,
+                response=copy.deepcopy(response),
+            )
+            self._append_event(invocation_id, attempt, {"type": "run.completed"})
+            return True
+
+    async def fail(self, invocation_id: str, attempt: int) -> bool:
+        async with self._lock:
+            state = self.states.get(invocation_id)
+            if state is None or not self._owns_attempt(state, attempt):
+                return False
+            self.states[invocation_id] = Invocation(
+                invocation_id=state.invocation_id,
+                status=InvocationStatus.FAILED,
+                attempt=state.attempt,
+                request=state.request,
+                response=None,
+            )
+            self._append_event(invocation_id, attempt, {"type": "run.failed"})
+            return True
+
+    async def append_event(
+        self,
+        invocation_id: str,
+        attempt: int,
+        event: JsonObject,
+    ) -> int | None:
+        async with self._lock:
+            state = self.states.get(invocation_id)
+            if not self._owns_attempt(state, attempt):
+                return None
+            persisted = self._append_event(invocation_id, attempt, event)
+            return persisted.sequence_number
+
+    async def events(
+        self,
+        invocation_id: str,
+        after_sequence: int | None = None,
+    ) -> list[InvocationEvent]:
+        _validate_invocation_id(invocation_id)
+        async with self._lock:
+            return [
+                copy.deepcopy(event)
+                for event in self.persisted_events
+                if event.invocation_id == invocation_id
+                and (after_sequence is None or event.sequence_number > after_sequence)
+            ]
+
+    @staticmethod
+    def _owns_attempt(state: Invocation | None, attempt: int) -> bool:
+        return bool(state and state.status == InvocationStatus.ACTIVE and state.attempt == attempt)
+
+    def _append_event(self, invocation_id: str, attempt: int, event: JsonObject) -> InvocationEvent:
+        persisted = InvocationEvent(
+            sequence_number=len(self.persisted_events) + 1,
+            invocation_id=invocation_id,
+            attempt=attempt,
+            event=copy.deepcopy(event),
+        )
+        self.persisted_events.append(persisted)
+        return persisted
+
+
+def runtime_store_is_persistent_environment() -> bool:
+    """Whether the process environment selects a complete durable Runtime Store."""
+    if os.getenv(RUNTIME_STORE_LOCAL_ENV, "").lower() == "true":
+        return False
+    managed = (
+        os.getenv(RUNTIME_STORE_LAKEBASE_BRANCH_ENV),
+        os.getenv(RUNTIME_STORE_DATABASE_ENV),
+        os.getenv(RUNTIME_STORE_USERNAME_ENV),
+    )
+    if any(managed):
+        return all(managed)
+    legacy = (
+        os.getenv(RUNTIME_STORE_LAKEBASE_ENDPOINT_ENV),
+        os.getenv(RUNTIME_STORE_SCHEMA_ENV),
+    )
+    return all(legacy)
+
+
+def runtime_store_from_environment() -> RuntimeStore:
+    """Construct the Runtime Store selected by the process environment.
+
+    The local marker overrides any inherited Lakebase variables. Managed Runtime Stores provide a
+    branch, database, and username; legacy Apps resources provide an endpoint plus ``PG*`` values.
+    An incomplete durable configuration is rejected rather than mixed with another resource.
+    """
+    if os.getenv(RUNTIME_STORE_LOCAL_ENV, "").lower() == "true":
+        return InMemoryRuntimeStore()
+
+    branch = os.getenv(RUNTIME_STORE_LAKEBASE_BRANCH_ENV)
+    database = os.getenv(RUNTIME_STORE_DATABASE_ENV)
+    username = os.getenv(RUNTIME_STORE_USERNAME_ENV)
+    managed = {
+        RUNTIME_STORE_LAKEBASE_BRANCH_ENV: branch,
+        RUNTIME_STORE_DATABASE_ENV: database,
+        RUNTIME_STORE_USERNAME_ENV: username,
+    }
+    if any(managed.values()):
+        missing = [name for name, value in managed.items() if not value]
+        if missing:
+            raise RuntimeError(
+                "Managed Runtime Store configuration is missing: " + ", ".join(missing)
+            )
+        from databricks_agentkit.runtime.durability.lakebase_runtime_store import (
+            LakebaseDurableRuntimeStore,
+        )
+
+        assert branch is not None
+        assert database is not None
+        assert username is not None
+        schema = os.getenv(RUNTIME_STORE_SCHEMA_ENV) or DEFAULT_RUNTIME_STORE_SCHEMA
+        return LakebaseDurableRuntimeStore.from_managed_runtime_store(
+            branch=branch,
+            database=database,
+            username=username,
+            schema=schema,
+        )
+
+    endpoint = os.getenv(RUNTIME_STORE_LAKEBASE_ENDPOINT_ENV)
+    schema = os.getenv(RUNTIME_STORE_SCHEMA_ENV)
+    if endpoint and schema:
+        from databricks_agentkit.runtime.durability.lakebase_runtime_store import (
+            LakebaseDurableRuntimeStore,
+        )
+
+        return LakebaseDurableRuntimeStore.from_app_resource(endpoint=endpoint, schema=schema)
+    if endpoint or schema:
+        raise RuntimeError(
+            f"{RUNTIME_STORE_LAKEBASE_ENDPOINT_ENV} and {RUNTIME_STORE_SCHEMA_ENV} must be set "
+            "together"
+        )
+    return InMemoryRuntimeStore()

databricks_agentkit/runtime/tool_manifest.py

@@ -0,0 +1,223 @@
+"""Read the framework-neutral tool bindings in the project's ``agent.toml``."""
+
+from __future__ import annotations
+
+import os
+import pathlib
+import re
+from collections.abc import Mapping
+from dataclasses import dataclass
+from typing import Any, cast
+
+try:
+    import tomllib  # ty: ignore[unresolved-import]
+except ModuleNotFoundError:
+    import tomli as tomllib
+
+# The env vars that carry the managed store bindings to the deployed runtime. (The agent.toml table
+# names live in `agent_project`, next to the manifest parsing that reads them.)
+MEMORY_STORE_ENV = "AGENT_MEMORY_STORE"
+SESSION_STORE_ENV = "AGENT_SESSION_STORE"
+
+
+class ToolManifestError(RuntimeError):
+    """Invalid declarative tool configuration that runtime adapters must surface."""
+
+
+@dataclass(frozen=True)
+class ScopeRecord:
+    kind: str
+    value: str
+    permission: str
+
+
+@dataclass(frozen=True)
+class ToolRecord:
+    id: str
+    kind: str
+    service: str | None = None
+    function: str | None = None
+    downscope: tuple[ScopeRecord, ...] = ()
+    space_id: str | None = None
+    auth: str | None = None
+
+
+def validate_genie_source(
+    tool_id: str, source: Mapping[str, Any], *, has_downscope: bool = False
+) -> None:
+    """Validate Genie bindings consistently for CLI models and direct manifest reads."""
+    kind = source.get("kind")
+    if kind not in {"genie_one", "genie_agent"}:
+        if "space_id" in source:
+            raise ToolManifestError("Only genie_agent bindings accept source.space_id.")
+        return
+    if not isinstance(tool_id, str) or not re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]{0,47}", tool_id):
+        raise ToolManifestError(
+            f"Invalid Genie tool id {tool_id!r}: use a function-name prefix of at most 48 "
+            "ASCII letters, digits, or underscores, starting with a letter or underscore."
+        )
+    allowed_fields = {"kind", "space_id"} if kind == "genie_agent" else {"kind"}
+    unexpected = source.keys() - allowed_fields
+    if unexpected:
+        raise ToolManifestError(
+            f"{kind} bindings do not accept source fields: {', '.join(sorted(unexpected))}."
+        )
+    if has_downscope:
+        raise ToolManifestError("Genie bindings do not accept policy.downscope.")
+    if kind == "genie_agent":
+        space_id = source.get("space_id")
+        if not isinstance(space_id, str) or not re.fullmatch(r"[0-9a-f]{32}", space_id):
+            raise ToolManifestError(
+                "Genie Agent space_id must be 32 lowercase hexadecimal characters."
+            )
+
+
+def project_root() -> pathlib.Path:
+    """Resolve the agent project containing ``agent.toml``.
+
+    This module ships in the databricks-agentbricks package, not inside the agent project, so locate the
+    project relative to where the agent runs (the current working directory) — not this file's
+    location. ``AGENTBRICKS_PROJECT_ROOT`` overrides for cases where the process starts elsewhere.
+    """
+    configured = os.getenv("AGENTBRICKS_PROJECT_ROOT")
+    if configured:
+        root = pathlib.Path(configured).expanduser().resolve()
+        if (root / "agent.toml").is_file():
+            return root
+        raise RuntimeError(f"AGENTBRICKS_PROJECT_ROOT has no agent.toml: {root}")
+
+    cwd = pathlib.Path.cwd().resolve()
+    for candidate in (cwd, *cwd.parents):
+        if (candidate / "agent.toml").is_file():
+            return candidate
+    raise RuntimeError(
+        "Could not locate agent.toml; set AGENTBRICKS_PROJECT_ROOT to the project root."
+    )
+
+
+def _required_string(value: object, description: str) -> str:
+    if not isinstance(value, str) or not value:
+        raise RuntimeError(f"agent.toml must declare {description}.")
+    return value
+
+
+def _scope(value: object) -> ScopeRecord:
+    if not isinstance(value, dict):
+        raise RuntimeError("agent.toml downscope entries must be tables.")
+    value = cast(dict[str, Any], value)
+    resource = _required_string(value.get("resource"), "a downscope resource")
+    kind, separator, resource_value = resource.partition(":")
+    if not separator or kind not in {"table", "volume", "workspace"} or not resource_value:
+        raise RuntimeError(f"Invalid agent.toml downscope resource: {resource!r}.")
+    permission = _required_string(value.get("permission", "read_only"), "a permission")
+    if permission not in {"read_only", "read_write"}:
+        raise RuntimeError(f"Invalid agent.toml downscope permission: {permission!r}.")
+    return ScopeRecord(kind=kind, value=resource_value, permission=permission)
+
+
+def _tool(value: object) -> ToolRecord:
+    if not isinstance(value, dict):
+        raise RuntimeError("agent.toml tools must be tables.")
+    value = cast(dict[str, Any], value)
+    source = value.get("source")
+    if not isinstance(source, dict):
+        raise RuntimeError("Each agent.toml tool must declare a source table.")
+    source = cast(dict[str, Any], source)
+    policy = value.get("policy", {})
+    if not isinstance(policy, dict):
+        raise RuntimeError("agent.toml tool policy must be a table.")
+    policy = cast(dict[str, Any], policy)
+    raw_downscope = policy.get("downscope", [])
+    if not isinstance(raw_downscope, list):
+        raise RuntimeError("agent.toml policy.downscope must be an array.")
+    kind = _required_string(source.get("kind"), "a tool source kind")
+    tool_id = _required_string(value.get("id"), "a tool id")
+    validate_genie_source(tool_id, source, has_downscope="downscope" in policy)
+    auth = value.get("auth")
+    if auth is not None and auth not in ("user", "app"):
+        raise ToolManifestError("Tool auth must be 'user' or 'app'.")
+    if kind == "uc_function" and auth == "user":
+        raise ToolManifestError("UC function auth supports only app/default identity.")
+    if kind == "python":
+        raise ToolManifestError(
+            "Python tools are code-first and cannot be declared in agent.toml. "
+            "Remove this entry; decorated tools in agent/tools remain active."
+        )
+    record = ToolRecord(
+        id=tool_id,
+        kind=kind,
+        service=source.get("service") if isinstance(source.get("service"), str) else None,
+        function=source.get("function") if isinstance(source.get("function"), str) else None,
+        space_id=source.get("space_id") if isinstance(source.get("space_id"), str) else None,
+        downscope=tuple(_scope(item) for item in raw_downscope),
+        auth=auth,
+    )
+    if record.kind == "sandbox" and (record.service != "system.ai.sandbox" or not record.downscope):
+        raise RuntimeError("Sandbox bindings require system.ai.sandbox and a downscope.")
+    if record.kind == "mcp" and not record.service:
+        raise RuntimeError("MCP bindings require source.service.")
+    if record.kind == "uc_function" and not record.function:
+        raise RuntimeError("UC function bindings require source.function.")
+    if record.kind not in {"sandbox", "mcp", "uc_function", "genie_one", "genie_agent"}:
+        raise RuntimeError(f"Unsupported agent.toml tool kind: {record.kind!r}.")
+    if record.kind != "sandbox" and record.downscope:
+        raise RuntimeError("Only sandbox bindings accept policy.downscope.")
+    return record
+
+
+def load_tools(*, expected_framework: str) -> tuple[ToolRecord, ...]:
+    """Load a fresh immutable view so direct manifest edits apply on the next request."""
+    path = project_root() / "agent.toml"
+    try:
+        with path.open("rb") as input_file:
+            document: dict[str, Any] = tomllib.load(input_file)
+    except (OSError, tomllib.TOMLDecodeError) as exc:
+        raise RuntimeError(f"Could not read {path}: {exc}") from exc
+    if document.get("schema_version") != 1:
+        raise RuntimeError(f"Unsupported agent.toml schema in {path}; expected schema_version = 1.")
+    agent = document.get("agent")
+    if not isinstance(agent, dict) or agent.get("framework") != expected_framework:
+        actual = agent.get("framework") if isinstance(agent, dict) else None
+        raise RuntimeError(
+            f"agent.toml framework {actual!r} does not match runtime {expected_framework!r}."
+        )
+    raw_tools = document.get("tools", [])
+    if not isinstance(raw_tools, list):
+        raise RuntimeError("agent.toml tools must be an array of tables.")
+    tools = tuple(_tool(item) for item in raw_tools)
+    ids = [tool.id for tool in tools]
+    if len(ids) != len(set(ids)):
+        raise RuntimeError("agent.toml tool ids must be unique.")
+    return tools
+
+
+def resolve_memory_store(explicit: str | None = None) -> str | None:
+    """The memory store id: ``explicit`` arg → ``AGENT_MEMORY_STORE`` env → None.
+
+    The one place the store-resolution precedence lives, shared by the framework adapters and the
+    chat-app UI so they always agree on which store is in effect. None means "no memory store".
+    """
+    return explicit or os.getenv(MEMORY_STORE_ENV) or None
+
+
+def resolve_session_store(explicit: str | None = None) -> str | None:
+    """The session store name: ``explicit`` arg → ``AGENT_SESSION_STORE`` env → None.
+
+    Shared precedence for the framework adapters and the chat-app UI (see ``resolve_memory_store``).
+    None means "no session store" (in-memory).
+    """
+    return explicit or os.getenv(SESSION_STORE_ENV) or None
+
+
+def downscope_wire(tool: ToolRecord) -> dict[str, list[dict[str, str]]]:
+    """Convert protected policy into the system.ai.sandbox MCP ``_meta`` shape."""
+    fields = {
+        "table": ("tables", "name"),
+        "volume": ("volumes", "name"),
+        "workspace": ("workspace_paths", "path"),
+    }
+    result: dict[str, list[dict[str, str]]] = {}
+    for scope in tool.downscope:
+        group, field = fields[scope.kind]
+        result.setdefault(group, []).append({field: scope.value, "permission": scope.permission})
+    return result

databricks_agentkit/runtime/tracing.py

@@ -0,0 +1,71 @@
+"""MLflow tracing setup — opt-in, enabled when MLflow has both a destination and an experiment.
+
+Tracing turns on only when a full config is present: a destination (``MLFLOW_TRACKING_URI`` or
+``MLFLOW_TRACING_DESTINATION``) AND an experiment (``MLFLOW_EXPERIMENT_ID`` or
+``MLFLOW_EXPERIMENT_NAME``) — whichever pair the user or the Apps resource binding provides. MLflow
+resolves the specific value itself; this only decides on/off. Requiring both halves avoids the
+half-configured case where traces silently export to a local file store instead of the workspace.
+When unconfigured, tracing is disabled outright so ``start_trace`` (which agents open around each
+invocation) has nothing to export to.
+"""
+
+import contextlib
+import os
+from collections.abc import Callable, Iterator
+from typing import Any
+
+import mlflow
+
+# Destination and experiment can each be named more than one way; accept any combination MLflow
+# understands (see mlflow.tracking.fluent._get_experiment_id_from_env for the experiment resolution).
+_DESTINATION_VARS = ("MLFLOW_TRACKING_URI", "MLFLOW_TRACING_DESTINATION")
+_EXPERIMENT_VARS = ("MLFLOW_EXPERIMENT_ID", "MLFLOW_EXPERIMENT_NAME")
+
+# Snapshotted once by configure_tracing() at startup (after .env is loaded) rather than at import, so
+# this module has no import-time side effects and load order does not matter.
+_enabled = False
+
+
+def configure_tracing(autolog: Callable[[], None] | None = None) -> None:
+    """Enable MLflow tracing for the agent. Call once at startup.
+
+    Reads the MLflow destination/experiment from the environment and no-ops (disables tracing) when
+    they are absent, so it is safe to call unconditionally. ``autolog`` is the framework's MLflow
+    autolog entry point (e.g. ``mlflow.langchain.autolog``), called only when tracing is enabled;
+    framework adapters bind it so callers get a zero-arg ``configure_tracing()``.
+    """
+    global _enabled
+    has_destination = any(os.getenv(v) for v in _DESTINATION_VARS)
+    has_experiment = any(os.getenv(v) for v in _EXPERIMENT_VARS)
+    _enabled = has_destination and has_experiment
+    if _enabled:
+        if autolog is not None:
+            autolog()
+    else:
+        # Agents wrap each invocation in start_trace; without an experiment it would try to export to a
+        # missing one (INVALID_PARAMETER_VALUE), so disable tracing outright.
+        mlflow.tracing.disable()
+
+
+@contextlib.contextmanager
+def start_trace(name: str, inputs: Any = None, session_id: str | None = None) -> Iterator[Any]:
+    """Start a trace by opening a root MLflow span around an operation (no-op when tracing is disabled).
+
+    Use this wherever tracing needs a root span to already exist. Framework autolog records spans
+    only when a trace is active and does not always start one itself (notably for streaming
+    invocations), so without a root span the operation can produce zero traces. Wrapping it in this
+    span gives autolog a parent to nest its spans under.
+
+    When ``session_id`` is given, the span's trace is tagged with it so traces group by session.
+
+    Yields the live span (so the caller can set outputs) when enabled, else ``None``.
+    """
+    if not _enabled:
+        yield None
+        return
+    with mlflow.start_span(name=name) as span:
+        if session_id:
+            mlflow.update_current_trace(metadata={"mlflow.trace.session": session_id})
+        if inputs is not None:
+            span.set_inputs(inputs)
+        yield span