coco-framework 0.1.0__py3-none-any.whl

coco_core/tracers/_context.py

@@ -0,0 +1,337 @@
+"""The tracing context and run record the framework uses in place of a vendor client.
+
+Upstream core imported these from its vendor's tracing client, which posts runs to a
+hosted service. The framework keeps the same names and shapes so upstream code runs
+unchanged, and nothing here leaves the process. Coco's record of what an agent did is
+the receipt the SDK writes, not a trace.
+
+`RunTree` keeps the fields, methods and printed form the tracers and their snapshot
+tests use, with the same defaults, so a run built here compares equal to the one
+upstream built. `coco_core.tracers.schemas.Run` is the same class.
+"""
+
+from __future__ import annotations
+
+import functools
+import os
+import secrets
+import time
+from collections.abc import Callable, Generator, Mapping, Sequence
+from contextlib import contextmanager
+from contextvars import ContextVar
+from datetime import datetime, timezone
+from typing import Any, TypeVar, cast
+from uuid import UUID
+
+from pydantic import BaseModel, ConfigDict, Field, model_validator
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+_PARENT_RUN_TREE: ContextVar[RunTree | None] = ContextVar("_PARENT_RUN_TREE", default=None)
+_PROJECT_NAME: ContextVar[str | None] = ContextVar("_PROJECT_NAME", default=None)
+_TAGS: ContextVar[list[str] | None] = ContextVar("_TAGS", default=None)
+_METADATA: ContextVar[dict[str, Any] | None] = ContextVar("_METADATA", default=None)
+_TRACING_ENABLED: ContextVar[bool | None] = ContextVar("_TRACING_ENABLED", default=None)
+_CLIENT: ContextVar[Any] = ContextVar("_CLIENT", default=None)
+_REPLICAS: ContextVar[Any] = ContextVar("_REPLICAS", default=None)
+_DISTRIBUTED_PARENT_ID: ContextVar[str | None] = ContextVar("_DISTRIBUTED_PARENT_ID", default=None)
+
+_CONTEXT_VARS: dict[str, ContextVar[Any]] = {
+    "parent": _PARENT_RUN_TREE,
+    "project_name": _PROJECT_NAME,
+    "tags": _TAGS,
+    "metadata": _METADATA,
+    "enabled": _TRACING_ENABLED,
+    "client": _CLIENT,
+    "replicas": _REPLICAS,
+    "distributed_parent_id": _DISTRIBUTED_PARENT_ID,
+}
+
+
+class Client:
+    """Stands where the vendor client stood in type hints. Never constructed here."""
+
+
+def get_tracing_context(context: Any = None) -> dict[str, Any]:
+    """The current tracing context: parent run, project, tags, metadata, enabled."""
+    if context is None:
+        return {key: var.get() for key, var in _CONTEXT_VARS.items()}
+    return {key: context.get(var) for key, var in _CONTEXT_VARS.items()}
+
+
+def _set_tracing_context(context: Mapping[str, Any] | None = None) -> None:
+    """Set the tracing context. `None` clears every key."""
+    if context is None:
+        for var in _CONTEXT_VARS.values():
+            var.set(None)
+        return
+    for key, value in context.items():
+        if key in _CONTEXT_VARS:
+            _CONTEXT_VARS[key].set(value)
+
+
+@contextmanager
+def tracing_context(**kwargs: Any) -> Generator[None, None, None]:
+    """Set tracing context keys for the duration of the block."""
+    previous = get_tracing_context()
+    _set_tracing_context({k: v for k, v in kwargs.items() if k in _CONTEXT_VARS})
+    try:
+        yield
+    finally:
+        _set_tracing_context(previous)
+
+
+def traceable(*args: Any, **kwargs: Any) -> Any:
+    """Accepts the vendor decorator's arguments and returns the function unchanged."""
+    del kwargs
+    if len(args) == 1 and callable(args[0]):
+        return args[0]
+
+    def decorator(func: F) -> F:
+        return func
+
+    return decorator
+
+
+@functools.lru_cache(maxsize=100)
+def get_env_var(
+    name: str,
+    default: str | None = None,
+    *,
+    namespaces: tuple[str, ...] = ("COCO",),
+) -> str | None:
+    """The first of `COCO_<name>` (by default) that is set, else `default`."""
+    for namespace in namespaces:
+        value = os.environ.get(f"{namespace}_{name}")
+        if value is not None:
+            return value
+    return default
+
+
+def tracing_is_enabled(ctx: Mapping[str, Any] | None = None) -> bool:
+    """Whether runs should be kept by a `CocoTracer`. Off unless asked for."""
+    context = ctx if ctx is not None else get_tracing_context()
+    if context.get("enabled") is not None:
+        return bool(context["enabled"])
+    return (get_env_var("TRACING_V2") or "").lower() == "true"
+
+
+def get_tracer_project(return_default_value: bool = True) -> str | None:  # noqa: FBT001, FBT002
+    """The project a run belongs to: `COCO_PROJECT`, else "default"."""
+    return os.environ.get("COCO_PROJECT") or ("default" if return_default_value else None)
+
+
+def get_coco_env_var_metadata() -> dict[str, str]:
+    """Environment variables the vendor client copied into run metadata. None here."""
+    return {}
+
+
+def uuid7() -> UUID:
+    """A time-ordered UUID (version 7), so runs sort by start."""
+    return _uuid7_from_ms(time.time_ns() // 1_000_000)
+
+
+def uuid7_from_datetime(dt: datetime) -> UUID:
+    """A version 7 UUID whose time part is `dt`."""
+    return _uuid7_from_ms(int(dt.timestamp() * 1000))
+
+
+def _uuid7_from_ms(ms: int) -> UUID:
+    rand = int.from_bytes(secrets.token_bytes(10), "big")
+    value = (ms & ((1 << 48) - 1)) << 80
+    value |= 0x7 << 76
+    value |= ((rand >> 62) & 0xFFF) << 64
+    value |= 0b10 << 62
+    value |= rand & ((1 << 62) - 1)
+    return UUID(int=value)
+
+
+def _dotted_order(start_time: datetime, run_id: UUID) -> str:
+    return start_time.strftime("%Y%m%dT%H%M%S%fZ") + str(run_id)
+
+
+class RunTree(BaseModel):
+    """One unit of work in a trace: a chain, model call, tool call or retriever call."""
+
+    id: UUID = Field(default_factory=uuid7)
+    name: str
+    start_time: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
+    run_type: str = "chain"
+    end_time: datetime | None = None
+    extra: dict = Field(default_factory=dict)
+    error: str | None = None
+    serialized: dict | None = None
+    events: list[dict] = Field(default_factory=list)
+    inputs: dict = Field(default_factory=dict)
+    outputs: dict | None = None
+    reference_example_id: UUID | None = None
+    parent_run_id: UUID | None = None
+    tags: list[str] | None = Field(default_factory=list)
+    attachments: dict = Field(default_factory=dict)
+    parent_run: RunTree | None = Field(default=None, exclude=True)
+    parent_dotted_order: str | None = Field(default=None, exclude=True)
+    child_runs: list[RunTree] = Field(default_factory=list, exclude=True)
+    session_name: str | None = Field(
+        default_factory=lambda: get_tracer_project() or "default", alias="project_name"
+    )
+    session_id: UUID | None = Field(default=None, alias="project_id")
+    ls_client: Any | None = Field(default=None, exclude=True)
+    dotted_order: str = ""
+    trace_id: UUID = Field(default=cast("UUID", ""))
+
+    model_config = ConfigDict(arbitrary_types_allowed=True, populate_by_name=True, extra="ignore")
+
+    @model_validator(mode="before")
+    @classmethod
+    def infer_defaults(cls, values: Any) -> Any:
+        """Fill the name, id, trace id and empty collections as upstream did."""
+        if not isinstance(values, dict):
+            return values
+        serialized = values.get("serialized")
+        if values.get("name") is None and serialized is not None:
+            if "name" in serialized:
+                values["name"] = serialized["name"]
+            elif "id" in serialized:
+                values["name"] = serialized["id"][-1]
+        if values.get("name") is None:
+            values["name"] = "Unnamed"
+        if "client" in values:
+            values["ls_client"] = values.pop("client")
+        elif "_client" in values:
+            values["ls_client"] = values.pop("_client")
+        parent_run = values.pop("parent_run", None)
+        if parent_run is not None:
+            values["parent_run_id"] = parent_run.id
+            values["parent_dotted_order"] = parent_run.dotted_order
+        if "id" not in values:
+            if values.get("start_time") is not None:
+                values["id"] = uuid7_from_datetime(values["start_time"])
+            else:
+                now = datetime.now(timezone.utc)
+                values["start_time"] = now
+                values["id"] = uuid7_from_datetime(now)
+        if "trace_id" not in values:
+            values["trace_id"] = parent_run.trace_id if parent_run is not None else values["id"]
+        values.setdefault("extra", {})
+        for key in ("events", "tags"):
+            if values.get(key) is None:
+                values[key] = []
+        if values.get("outputs") is None:
+            values["outputs"] = {}
+        for key in ("inputs", "outputs"):
+            value = values.get(key)
+            if value is not None and not isinstance(value, dict):
+                values[key] = {"output": value}
+        if values.get("attachments") is None:
+            values["attachments"] = {}
+        return values
+
+    @model_validator(mode="after")
+    def ensure_dotted_order(self) -> RunTree:
+        """Give the run its place in the tree when nothing set one."""
+        if self.dotted_order and self.dotted_order.strip():
+            return self
+        current = _dotted_order(self.start_time, self.id)
+        if self.parent_dotted_order is not None:
+            self.dotted_order = self.parent_dotted_order + "." + current
+        else:
+            self.dotted_order = current
+        return self
+
+    @property
+    def client(self) -> Any:
+        """Whatever the caller attached. Nothing is created here."""
+        return self.ls_client
+
+    @property
+    def metadata(self) -> dict[str, Any]:
+        """The run's metadata, stored under `extra["metadata"]`."""
+        if self.extra is None:
+            self.extra = {}
+        return self.extra.setdefault("metadata", {})
+
+    @property
+    def revision_id(self) -> UUID | None:
+        """The revision id from metadata, if any."""
+        return self.metadata.get("revision_id")
+
+    @property
+    def latency(self) -> float | None:
+        """Seconds from start to end, once ended."""
+        if self.end_time is None:
+            return None
+        return (self.end_time - self.start_time).total_seconds()
+
+    def add_tags(self, tags: Sequence[str] | str) -> None:
+        """Add tags to the run."""
+        if isinstance(tags, str):
+            tags = [tags]
+        if self.tags is None:
+            self.tags = []
+        self.tags.extend(tags)
+
+    def add_metadata(self, metadata: dict[str, Any]) -> None:
+        """Merge metadata into the run."""
+        self.metadata.update(metadata)
+
+    def add_outputs(self, outputs: dict[str, Any]) -> None:
+        """Merge outputs into the run."""
+        if self.outputs is None:
+            self.outputs = {}
+        self.outputs.update(outputs)
+
+    def add_inputs(self, inputs: dict[str, Any]) -> None:
+        """Merge inputs into the run."""
+        if self.inputs is None:
+            self.inputs = {}
+        self.inputs.update(inputs)
+
+    def add_event(self, events: Sequence[dict] | dict | str) -> None:
+        """Append one event, several, or a message recorded as an event."""
+        if self.events is None:
+            self.events = []
+        if isinstance(events, dict):
+            self.events.append(events)
+        elif isinstance(events, str):
+            self.events.append(
+                {"name": "event", "time": datetime.now(timezone.utc).isoformat(), "message": events}
+            )
+        else:
+            self.events.extend(events)
+
+    def end(
+        self,
+        *,
+        outputs: dict | None = None,
+        error: str | None = None,
+        end_time: datetime | None = None,
+        events: Sequence[dict] | None = None,
+        metadata: dict[str, Any] | None = None,
+    ) -> None:
+        """Close the run."""
+        self.end_time = end_time or datetime.now(timezone.utc)
+        if outputs is not None:
+            self.add_outputs(outputs)
+        if error is not None:
+            self.error = error
+        if events is not None:
+            self.add_event(list(events))
+        if metadata is not None:
+            self.add_metadata(metadata)
+
+    def create_child(self, name: str, run_type: str = "chain", **kwargs: Any) -> RunTree:
+        """A child run under this one."""
+        child = RunTree(name=name, run_type=run_type, parent_run=self, **kwargs)
+        self.child_runs.append(child)
+        return child
+
+    def __repr__(self) -> str:
+        return (
+            f"RunTree(id={self.id}, name='{self.name}', "
+            f"run_type='{self.run_type}', dotted_order='{self.dotted_order}')"
+        )
+
+
+RunTree.model_rebuild()
+
+Run = RunTree

coco_core/tracers/_streaming.py

@@ -0,0 +1,53 @@
+"""Internal tracers used for `stream_log` and `astream` events implementations."""
+
+import typing
+from collections.abc import AsyncIterator, Iterator
+from uuid import UUID
+
+T = typing.TypeVar("T")
+
+
+# THIS IS USED IN COCO_GRAPH.
+@typing.runtime_checkable
+class _StreamingCallbackHandler(typing.Protocol[T]):
+    """Types for streaming callback handlers.
+
+    This is a common mixin that the callback handlers for both astream events and
+    astream log inherit from.
+
+    The `tap_output_aiter` method is invoked in some contexts to produce callbacks for
+    intermediate results.
+    """
+
+    def tap_output_aiter(
+        self, run_id: UUID, output: AsyncIterator[T]
+    ) -> AsyncIterator[T]:
+        """Used for internal astream_log and astream events implementations."""
+
+    def tap_output_iter(self, run_id: UUID, output: Iterator[T]) -> Iterator[T]:
+        """Used for internal astream_log and astream events implementations."""
+
+
+# THIS IS USED IN COCO_GRAPH.
+class _V2StreamingCallbackHandler:
+    """Marker base class for handlers that consume `on_stream_event` (v2).
+
+    A handler inheriting from this class signals that it wants content-
+    block lifecycle events from `stream_events(version="v3")` (and its
+    async equivalent) rather than the v1 `on_llm_new_token` chunks.
+    `BaseChatModel.invoke` uses
+    `isinstance(handler, _V2StreamingCallbackHandler)` to decide whether
+    to route an invoke through the v2 event generator.
+
+    Implemented as a concrete marker class (not a `Protocol`) so opt-in
+    is explicit via inheritance. An empty `runtime_checkable` Protocol
+    would match every object and misroute every call. The event
+    delivery contract itself lives on
+    `BaseCallbackHandler.on_stream_event`.
+    """
+
+
+__all__ = [
+    "_StreamingCallbackHandler",
+    "_V2StreamingCallbackHandler",
+]