coco-framework 0.1.0__py3-none-any.whl

coco_core/tracers/coco.py

@@ -0,0 +1,279 @@
+"""A tracer that keeps runs in memory. It sends nothing anywhere.
+
+Upstream's tracer at this path posted every run to its vendor's hosted service. Coco
+keeps the class, so the callback manager, the graph runtime and user code that look
+for it still work, and drops the network. The last finished root run stays on
+`latest_run`, which is where the Coco receipt reads from.
+"""
+
+from __future__ import annotations
+
+import logging
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING, Any, cast
+from uuid import UUID
+
+from typing_extensions import override
+
+from coco_core.env import get_runtime_environment
+from coco_core.load import dumpd
+from coco_core.messages.ai import UsageMetadata, add_usage
+from coco_core.tracers._compat import run_construct, run_to_dict
+from coco_core.tracers._context import get_tracer_project, get_tracing_context
+from coco_core.tracers.base import BaseTracer
+from coco_core.tracers.schemas import Run
+
+if TYPE_CHECKING:
+    from collections.abc import Mapping
+
+    from coco_core.messages import BaseMessage
+    from coco_core.outputs import ChatGenerationChunk, GenerationChunk
+
+logger = logging.getLogger(__name__)
+_LOGGED: set[tuple[str, type[Exception]]] = set()
+
+OVERRIDABLE_LANGSMITH_INHERITABLE_METADATA_KEYS: frozenset[str] = frozenset({"ls_agent_type"})
+"""Tracing metadata keys a nested caller may override. Every other key is first wins."""
+
+
+def log_error_once(method: str, exception: Exception) -> None:
+    """Log an error once per method and exception type."""
+    if (method, type(exception)) in _LOGGED:
+        return
+    _LOGGED.add((method, type(exception)))
+    logger.error(exception)
+
+
+def wait_for_all_tracers() -> None:
+    """Nothing is queued, so there is nothing to wait for."""
+
+
+def _get_usage_metadata_from_generations(
+    generations: list[list[dict[str, Any]]],
+) -> UsageMetadata | None:
+    """Sum `usage_metadata` across serialized generations."""
+    output: UsageMetadata | None = None
+    for batch in generations:
+        for generation in batch:
+            if isinstance(generation, dict) and "message" in generation:
+                usage = _get_usage_metadata_from_message(generation["message"])
+                if usage is not None:
+                    output = add_usage(output, usage)
+    return output
+
+
+def _get_usage_metadata_from_message(message: Any) -> UsageMetadata | None:
+    if not isinstance(message, dict):
+        return None
+    kwargs = message.get("kwargs")
+    if isinstance(kwargs, dict) and isinstance(kwargs.get("usage_metadata"), dict):
+        return cast("UsageMetadata", kwargs["usage_metadata"])
+    return None
+
+
+class CocoTracer(BaseTracer):
+    """Keeps each trace's runs in memory, with project, tags and metadata applied."""
+
+    run_inline = True
+
+    def __init__(
+        self,
+        example_id: UUID | str | None = None,
+        project_name: str | None = None,
+        client: Any | None = None,
+        tags: list[str] | None = None,
+        *,
+        metadata: Mapping[str, str] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """Create the tracer.
+
+        Args:
+            example_id: Attached to root runs as `reference_example_id`.
+            project_name: The project on each run. Defaults to `COCO_PROJECT`.
+            client: Accepted for signature compatibility and never called.
+            tags: Added to every run.
+            metadata: Added to a run's metadata where the run lacks the key.
+            **kwargs: Passed to `BaseTracer`.
+        """
+        super().__init__(**kwargs)
+        self.example_id = UUID(example_id) if isinstance(example_id, str) else example_id
+        self.project_name = project_name or get_tracer_project()
+        self.client = client
+        self.tags = tags or []
+        self.latest_run: Run | None = None
+        self.run_has_token_event_map: dict[str, bool] = {}
+        self.tracing_metadata: dict[str, str] | None = (
+            dict(metadata) if metadata is not None else None
+        )
+
+    def copy_with_metadata_defaults(
+        self,
+        *,
+        metadata: Mapping[str, str] | None = None,
+        tags: list[str] | None = None,
+    ) -> CocoTracer:
+        """Return a new tracer with merged tracer-only defaults."""
+        base = self.tracing_metadata
+        if metadata is None:
+            merged = dict(base) if base is not None else None
+        elif base is None:
+            merged = dict(metadata)
+        else:
+            merged = dict(base)
+            for key, value in metadata.items():
+                if key not in merged or key in OVERRIDABLE_LANGSMITH_INHERITABLE_METADATA_KEYS:
+                    merged[key] = value
+        merged_tags = sorted(set(self.tags + tags)) if tags else self.tags
+        return self.__class__(
+            example_id=self.example_id,
+            project_name=self.project_name,
+            client=self.client,
+            tags=merged_tags,
+            metadata=merged,
+            run_map=self.run_map,
+            order_map=self.order_map,
+            _external_run_ids=self._external_run_ids,
+        )
+
+    def _start_trace(self, run: Run) -> None:
+        if self.project_name:
+            run.session_name = self.project_name
+        if self.tags is not None:
+            run.tags = sorted(set(run.tags + self.tags)) if run.tags else self.tags.copy()
+        super()._start_trace(run)
+        if get_tracing_context().get("enabled") is False:
+            run.extra["__disabled"] = True
+
+    def on_chat_model_start(
+        self,
+        serialized: dict[str, Any],
+        messages: list[list[BaseMessage]],
+        *,
+        run_id: UUID,
+        tags: list[str] | None = None,
+        parent_run_id: UUID | None = None,
+        metadata: dict[str, Any] | None = None,
+        name: str | None = None,
+        **kwargs: Any,
+    ) -> Run:
+        """Start a run for a chat model call, recorded as an LLM run."""
+        start_time = datetime.now(timezone.utc)
+        if metadata:
+            kwargs.update({"metadata": metadata})
+        run = Run(
+            id=run_id,
+            parent_run_id=parent_run_id,
+            serialized=serialized,
+            inputs={"messages": [[dumpd(msg) for msg in batch] for batch in messages]},
+            extra=kwargs,
+            events=[{"name": "start", "time": start_time}],
+            start_time=start_time,
+            run_type="llm",
+            tags=tags,
+            name=name,
+        )
+        self._start_trace(run)
+        self._on_chat_model_start(run)
+        return run
+
+    def _persist_run(self, run: Run) -> None:
+        # Keep the finished root without its whole child tree.
+        data = run_to_dict(run, exclude={"child_runs", "inputs", "outputs"})
+        self.latest_run = run_construct(**data, inputs=run.inputs, outputs=run.outputs)
+
+    def get_run_url(self) -> str:
+        """Traces stay in this process, so there is no URL to give.
+
+        Raises:
+            ValueError: Always.
+        """
+        msg = "CocoTracer keeps runs in memory; there is no run URL."
+        raise ValueError(msg)
+
+    def _get_tags(self, run: Run) -> list[str]:
+        tags = set(run.tags or [])
+        tags.update(self.tags or [])
+        return list(tags)
+
+    def _record(self, run: Run) -> None:
+        if run.extra.get("__disabled"):
+            return
+        run.extra["runtime"] = get_runtime_environment()
+        run.tags = self._get_tags(run)
+        _patch_missing_metadata(self, run)
+
+    def _on_llm_start(self, run: Run) -> None:
+        if run.parent_run_id is None:
+            run.reference_example_id = self.example_id
+        self._record(run)
+
+    @override
+    def _llm_run_with_token_event(
+        self,
+        token: str | list[str | dict[str, Any]],
+        run_id: UUID,
+        chunk: GenerationChunk | ChatGenerationChunk | None = None,
+        parent_run_id: UUID | None = None,
+    ) -> Run:
+        run_id_str = str(run_id)
+        if run_id_str not in self.run_has_token_event_map:
+            self.run_has_token_event_map[run_id_str] = True
+        else:
+            return self._get_run(run_id, run_type={"llm", "chat_model"})
+        return super()._llm_run_with_token_event(
+            token, run_id, chunk=None, parent_run_id=parent_run_id
+        )
+
+    def _on_chat_model_start(self, run: Run) -> None:
+        if run.parent_run_id is None:
+            run.reference_example_id = self.example_id
+        self._record(run)
+
+    def _on_llm_end(self, run: Run) -> None:
+        if run.outputs and "generations" in run.outputs:
+            usage = _get_usage_metadata_from_generations(run.outputs["generations"])
+            if usage is not None:
+                run.extra.setdefault("metadata", {})["usage_metadata"] = usage
+
+    def _on_chain_start(self, run: Run) -> None:
+        if run.parent_run_id is None:
+            run.reference_example_id = self.example_id
+        if not run.extra.get("defers_inputs"):
+            self._record(run)
+
+    def _on_chain_end(self, run: Run) -> None:
+        if run.extra.get("defers_inputs"):
+            self._record(run)
+
+    def _on_chain_error(self, run: Run) -> None:
+        if run.extra.get("defers_inputs"):
+            self._record(run)
+
+    def _on_tool_start(self, run: Run) -> None:
+        if run.parent_run_id is None:
+            run.reference_example_id = self.example_id
+        self._record(run)
+
+    def _on_retriever_start(self, run: Run) -> None:
+        if run.parent_run_id is None:
+            run.reference_example_id = self.example_id
+        self._record(run)
+
+    def wait_for_futures(self) -> None:
+        """Nothing is queued, so there is nothing to wait for."""
+
+
+def _patch_missing_metadata(self: CocoTracer, run: Run) -> None:
+    if not self.tracing_metadata:
+        return
+    metadata = run.metadata
+    patched = None
+    for k, v in self.tracing_metadata.items():
+        if k not in metadata or k in OVERRIDABLE_LANGSMITH_INHERITABLE_METADATA_KEYS:
+            if k in metadata and metadata[k] == v:
+                continue
+            if patched is None:
+                patched = {**metadata}
+                run.extra["metadata"] = patched
+            patched[k] = v

coco_core/tracers/context.py

@@ -0,0 +1,205 @@
+"""Context management for tracers."""
+
+from __future__ import annotations
+
+from contextlib import contextmanager
+from contextvars import ContextVar
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Literal,
+    cast,
+)
+from uuid import UUID
+
+from coco_core.tracers import _context as ls_rh
+from coco_core.tracers import _context as ls_utils
+
+from coco_core.tracers.coco import CocoTracer
+from coco_core.tracers.run_collector import RunCollectorCallbackHandler
+
+if TYPE_CHECKING:
+    from collections.abc import Generator
+
+    from coco_core.tracers._context import Client as LangSmithClient
+
+    from coco_core.callbacks.base import BaseCallbackHandler, Callbacks
+    from coco_core.callbacks.manager import AsyncCallbackManager, CallbackManager
+
+# for backwards partial compatibility if this is imported by users but unused
+tracing_callback_var: Any = None
+tracing_v2_callback_var: ContextVar[CocoTracer | None] = ContextVar(
+    "tracing_callback_v2", default=None
+)
+run_collector_var: ContextVar[RunCollectorCallbackHandler | None] = ContextVar(
+    "run_collector", default=None
+)
+
+
+@contextmanager
+def tracing_v2_enabled(
+    project_name: str | None = None,
+    *,
+    example_id: str | UUID | None = None,
+    tags: list[str] | None = None,
+    client: LangSmithClient | None = None,
+) -> Generator[CocoTracer, None, None]:
+    """Instruct Coco to log all runs in context to LangSmith.
+
+    Args:
+        project_name: The name of the project.
+
+            Defaults to `'default'`.
+        example_id: The ID of the example.
+        tags: The tags to add to the run.
+        client: The client of the langsmith.
+
+    Yields:
+        The Coco tracer.
+
+    Example:
+        >>> with tracing_v2_enabled():
+        ...     # Coco code will automatically be traced
+
+        You can use this to fetch the LangSmith run URL:
+
+        >>> with tracing_v2_enabled() as cb:
+        ...     chain.invoke("foo")
+        ...     run_url = cb.get_run_url()
+    """
+    if isinstance(example_id, str):
+        example_id = UUID(example_id)
+    cb = CocoTracer(
+        example_id=example_id,
+        project_name=project_name,
+        tags=tags,
+        client=client,
+    )
+    token = tracing_v2_callback_var.set(cb)
+    try:
+        yield cb
+    finally:
+        tracing_v2_callback_var.reset(token)
+
+
+@contextmanager
+def collect_runs() -> Generator[RunCollectorCallbackHandler, None, None]:
+    """Collect all run traces in context.
+
+    Yields:
+        The run collector callback handler.
+
+    Example:
+        >>> with collect_runs() as runs_cb:
+                chain.invoke("foo")
+                run_id = runs_cb.traced_runs[0].id
+    """
+    cb = RunCollectorCallbackHandler()
+    token = run_collector_var.set(cb)
+    try:
+        yield cb
+    finally:
+        run_collector_var.reset(token)
+
+
+def _get_trace_callbacks(
+    project_name: str | None = None,
+    example_id: str | UUID | None = None,
+    callback_manager: CallbackManager | AsyncCallbackManager | None = None,
+) -> Callbacks:
+    if _tracing_v2_is_enabled():
+        project_name_ = project_name or _get_tracer_project()
+        tracer = tracing_v2_callback_var.get() or CocoTracer(
+            project_name=project_name_,
+            example_id=example_id,
+        )
+        if callback_manager is None:
+            cb = cast("Callbacks", [tracer])
+        else:
+            if not any(
+                isinstance(handler, CocoTracer)
+                for handler in callback_manager.handlers
+            ):
+                callback_manager.add_handler(tracer)
+                # If it already has a CocoTracer, we don't need to add another one.
+                # this would likely mess up the trace hierarchy.
+            cb = callback_manager
+    else:
+        cb = None
+    return cb
+
+
+def _tracing_v2_is_enabled() -> bool | Literal["local"]:
+    if tracing_v2_callback_var.get() is not None:
+        return True
+    return ls_utils.tracing_is_enabled()
+
+
+def _get_tracer_project() -> str:
+    tracing_context = ls_rh.get_tracing_context()
+    run_tree = tracing_context["parent"]
+    if run_tree is None and tracing_context["project_name"] is not None:
+        return cast("str", tracing_context["project_name"])
+    return getattr(
+        run_tree,
+        "session_name",
+        getattr(
+            # Note, if people are trying to nest @traceable functions and the
+            # tracing_v2_enabled context manager, this will likely mess up the
+            # tree structure.
+            tracing_v2_callback_var.get(),
+            "project",
+            # Have to set this to a string even though it always will return
+            # a string because `get_tracer_project` technically can return
+            # None, but only when a specific argument is supplied.
+            # Therefore, this just tricks the mypy type checker
+            str(ls_utils.get_tracer_project()),
+        ),
+    )
+
+
+_configure_hooks: list[
+    tuple[
+        ContextVar[BaseCallbackHandler | None],
+        bool,
+        type[BaseCallbackHandler] | None,
+        str | None,
+    ]
+] = []
+
+
+def register_configure_hook(
+    context_var: ContextVar[Any | None],
+    inheritable: bool,  # noqa: FBT001
+    handle_class: type[BaseCallbackHandler] | None = None,
+    env_var: str | None = None,
+) -> None:
+    """Register a configure hook.
+
+    Args:
+        context_var: The context variable.
+        inheritable: Whether the context variable is inheritable.
+        handle_class: The callback handler class.
+        env_var: The environment variable.
+
+    Raises:
+        ValueError: If `env_var` is set, `handle_class` must also be set to a non-`None`
+            value.
+    """
+    if env_var is not None and handle_class is None:
+        msg = "If env_var is set, handle_class must also be set to a non-None value."
+        raise ValueError(msg)
+
+    _configure_hooks.append(
+        (
+            # the typings of ContextVar do not have the generic arg set as covariant
+            # so we have to cast it
+            cast("ContextVar[BaseCallbackHandler | None]", context_var),
+            inheritable,
+            handle_class,
+            env_var,
+        )
+    )
+
+
+register_configure_hook(run_collector_var, inheritable=False)