codekingpro/portable-devtools

1"""A tracer implementation that records to LangChain endpoint."""2 3from __future__ import annotations4 5import logging6from concurrent.futures import ThreadPoolExecutor7from datetime import datetime, timezone8from typing import TYPE_CHECKING, Any, cast9from uuid import UUID10 11from langsmith import Client, get_tracing_context12from langsmith import run_trees as rt13from langsmith import utils as ls_utils14from tenacity import (15    Retrying,16    retry_if_exception_type,17    stop_after_attempt,18    wait_exponential_jitter,19)20from typing_extensions import override21 22from langchain_core.env import get_runtime_environment23from langchain_core.load import dumpd24from langchain_core.messages.ai import UsageMetadata, add_usage25from langchain_core.tracers._compat import run_construct, run_to_dict26from langchain_core.tracers.base import BaseTracer27from langchain_core.tracers.schemas import Run28 29if TYPE_CHECKING:30    from collections.abc import Mapping31 32    from langchain_core.messages import BaseMessage33    from langchain_core.outputs import ChatGenerationChunk, GenerationChunk34 35logger = logging.getLogger(__name__)36_LOGGED = set()37_EXECUTOR: ThreadPoolExecutor | None = None38 39OVERRIDABLE_LANGSMITH_INHERITABLE_METADATA_KEYS: frozenset[str] = frozenset(40    {"ls_agent_type"}41)42"""Allowlist of LangSmith-only tracing metadata keys that bypass the default43"first wins" merge semantics used when propagating tracer metadata to nested44runs.45 46Keys in this set are ALWAYS overridden by the nearest enclosing tracer config,47so nested callers (e.g. a subagent) can replace a value inherited from an48ancestor.49 50Keep this list very small: every key here loses the default "first wins"51protection and is always clobbered by the nearest enclosing tracer config.52Only keys that are strictly for LangSmith tracing bookkeeping should be added.53"""54 55 56def log_error_once(method: str, exception: Exception) -> None:57    """Log an error once.58 59    Args:60        method: The method that raised the exception.61        exception: The exception that was raised.62    """63    if (method, type(exception)) in _LOGGED:64        return65    _LOGGED.add((method, type(exception)))66    logger.error(exception)67 68 69def wait_for_all_tracers() -> None:70    """Wait for all tracers to finish."""71    if rt._CLIENT is not None:  # noqa: SLF00172        rt._CLIENT.flush()  # noqa: SLF00173 74 75def get_client() -> Client:76    """Get the client.77 78    Returns:79        The LangSmith client.80    """81    return rt.get_cached_client()82 83 84def _get_executor() -> ThreadPoolExecutor:85    """Get the executor."""86    global _EXECUTOR  # noqa: PLW060387    if _EXECUTOR is None:88        _EXECUTOR = ThreadPoolExecutor()89    return _EXECUTOR90 91 92def _get_usage_metadata_from_generations(93    generations: list[list[dict[str, Any]]],94) -> UsageMetadata | None:95    """Extract and aggregate `usage_metadata` from generations.96 97    Iterates through generations to find and aggregate all `usage_metadata` found in98    messages. This expects the serialized message payload shape produced by tracer99    internals:100 101        `{"message": {"kwargs": {"usage_metadata": {...}}}}`102 103    Args:104        generations: List of generation batches, where each batch is a list of105            generation dicts that may contain a `'message'` key with106            usage metadata.107 108    Returns:109        The aggregated `usage_metadata` dict if found, otherwise `None`.110    """111    output: UsageMetadata | None = None112    for generation_batch in generations:113        for generation in generation_batch:114            if isinstance(generation, dict) and "message" in generation:115                message = generation["message"]116                usage_metadata = _get_usage_metadata_from_message(message)117                if usage_metadata is not None:118                    output = add_usage(output, usage_metadata)119    return output120 121 122def _get_usage_metadata_from_message(message: Any) -> UsageMetadata | None:123    """Extract usage metadata from a generation's message payload."""124    if not isinstance(message, dict):125        return None126 127    kwargs = message.get("kwargs")128    if isinstance(kwargs, dict) and isinstance(kwargs.get("usage_metadata"), dict):129        return cast("UsageMetadata", kwargs["usage_metadata"])130 131    return None132 133 134class LangChainTracer(BaseTracer):135    """Implementation of the `SharedTracer` that `POSTS` to the LangChain endpoint."""136 137    run_inline = True138 139    def __init__(140        self,141        example_id: UUID | str | None = None,142        project_name: str | None = None,143        client: Client | None = None,144        tags: list[str] | None = None,145        *,146        metadata: Mapping[str, str] | None = None,147        **kwargs: Any,148    ) -> None:149        """Initialize the LangChain tracer.150 151        Args:152            example_id: The example ID.153            project_name: The project name.154 155                Defaults to the tracer project.156            client: The client.157 158                Defaults to the global client.159            tags: The tags.160 161                Defaults to an empty list.162            metadata: Additional metadata to include if it isn't already in the run.163 164                Defaults to None.165            **kwargs: Additional keyword arguments.166        """167        super().__init__(**kwargs)168        self.example_id = (169            UUID(example_id) if isinstance(example_id, str) else example_id170        )171        self.project_name = project_name or ls_utils.get_tracer_project()172        self.client = client or get_client()173        self.tags = tags or []174        self.latest_run: Run | None = None175        self.run_has_token_event_map: dict[str, bool] = {}176        self.tracing_metadata: dict[str, str] | None = (177            dict(metadata) if metadata is not None else None178        )179 180    def copy_with_metadata_defaults(181        self,182        *,183        metadata: Mapping[str, str] | None = None,184        tags: list[str] | None = None,185    ) -> LangChainTracer:186        """Return a new tracer with merged tracer-only defaults."""187        base_metadata = self.tracing_metadata188        if metadata is None:189            merged_metadata = dict(base_metadata) if base_metadata is not None else None190        elif base_metadata is None:191            merged_metadata = dict(metadata)192        else:193            merged_metadata = dict(base_metadata)194            for key, value in metadata.items():195                # For allowlisted LangSmith-only inheritable metadata keys196                # (e.g. ``ls_agent_type``), nested callers are allowed to197                # OVERRIDE the value inherited from an ancestor. For all198                # other keys we keep the existing "first wins" behavior so199                # that ancestor-provided tracing metadata is not accidentally200                # clobbered by child runs.201                if (202                    key not in merged_metadata203                    or key in OVERRIDABLE_LANGSMITH_INHERITABLE_METADATA_KEYS204                ):205                    merged_metadata[key] = value206 207        merged_tags = sorted(set(self.tags + tags)) if tags else self.tags208 209        return self.__class__(210            example_id=self.example_id,211            project_name=self.project_name,212            client=self.client,213            tags=merged_tags,214            metadata=merged_metadata,215            run_map=self.run_map,216            order_map=self.order_map,217            _external_run_ids=self._external_run_ids,218        )219 220    def _start_trace(self, run: Run) -> None:221        if self.project_name:222            run.session_name = self.project_name223        if self.tags is not None:224            if run.tags:225                run.tags = sorted(set(run.tags + self.tags))226            else:227                run.tags = self.tags.copy()228 229        super()._start_trace(run)230        if run.ls_client is None:231            run.ls_client = self.client232        if get_tracing_context().get("enabled") is False:233            run.extra["__disabled"] = True234 235    def on_chat_model_start(236        self,237        serialized: dict[str, Any],238        messages: list[list[BaseMessage]],239        *,240        run_id: UUID,241        tags: list[str] | None = None,242        parent_run_id: UUID | None = None,243        metadata: dict[str, Any] | None = None,244        name: str | None = None,245        **kwargs: Any,246    ) -> Run:247        """Start a trace for an LLM run.248 249        Args:250            serialized: The serialized model.251            messages: The messages.252            run_id: The run ID.253            tags: The tags.254            parent_run_id: The parent run ID.255            metadata: The metadata.256            name: The name.257            **kwargs: Additional keyword arguments.258 259        Returns:260            The run.261        """262        start_time = datetime.now(timezone.utc)263        if metadata:264            kwargs.update({"metadata": metadata})265        chat_model_run = Run(266            id=run_id,267            parent_run_id=parent_run_id,268            serialized=serialized,269            inputs={"messages": [[dumpd(msg) for msg in batch] for batch in messages]},270            extra=kwargs,271            events=[{"name": "start", "time": start_time}],272            start_time=start_time,273            run_type="llm",274            tags=tags,275            name=name,276        )277        self._start_trace(chat_model_run)278        self._on_chat_model_start(chat_model_run)279        return chat_model_run280 281    def _persist_run(self, run: Run) -> None:282        # We want to free up more memory by avoiding keeping a reference to the283        # whole nested run tree.284        run_data = run_to_dict(run, exclude={"child_runs", "inputs", "outputs"})285        self.latest_run = run_construct(286            **run_data,287            inputs=run.inputs,288            outputs=run.outputs,289        )290 291    def get_run_url(self) -> str:292        """Get the LangSmith root run URL.293 294        Returns:295            The LangSmith root run URL.296 297        Raises:298            ValueError: If no traced run is found.299            ValueError: If the run URL cannot be found.300        """301        if not self.latest_run:302            msg = "No traced run found."303            raise ValueError(msg)304        # If this is the first run in a project, the project may not yet be created.305        # This method is only really useful for debugging flows, so we will assume306        # there is some tolerace for latency.307        for attempt in Retrying(308            stop=stop_after_attempt(5),309            wait=wait_exponential_jitter(),310            retry=retry_if_exception_type(ls_utils.LangSmithError),311        ):312            with attempt:313                return self.client.get_run_url(314                    run=self.latest_run, project_name=self.project_name315                )316        msg = "Failed to get run URL."317        raise ValueError(msg)318 319    def _get_tags(self, run: Run) -> list[str]:320        """Get combined tags for a run."""321        tags = set(run.tags or [])322        tags.update(self.tags or [])323        return list(tags)324 325    def _persist_run_single(self, run: Run) -> None:326        """Persist a run."""327        if run.extra.get("__disabled"):328            return329        try:330            run.extra["runtime"] = get_runtime_environment()331            run.tags = self._get_tags(run)332            _patch_missing_metadata(self, run)333            if run.ls_client is not self.client:334                run.ls_client = self.client335            run.post()336        except Exception as e:337            # Errors are swallowed by the thread executor so we need to log them here338            log_error_once("post", e)339            raise340 341    @staticmethod342    def _update_run_single(run: Run) -> None:343        """Update a run."""344        if run.extra.get("__disabled"):345            return346        try:347            run.patch(exclude_inputs=run.extra.get("inputs_is_truthy", False))348        except Exception as e:349            # Errors are swallowed by the thread executor so we need to log them here350            log_error_once("patch", e)351            raise352 353    def _on_llm_start(self, run: Run) -> None:354        """Persist an LLM run."""355        if run.parent_run_id is None:356            run.reference_example_id = self.example_id357        self._persist_run_single(run)358 359    @override360    def _llm_run_with_token_event(361        self,362        token: str,363        run_id: UUID,364        chunk: GenerationChunk | ChatGenerationChunk | None = None,365        parent_run_id: UUID | None = None,366    ) -> Run:367        run_id_str = str(run_id)368        if run_id_str not in self.run_has_token_event_map:369            self.run_has_token_event_map[run_id_str] = True370        else:371            return self._get_run(run_id, run_type={"llm", "chat_model"})372        return super()._llm_run_with_token_event(373            # Drop the chunk; we don't need to save it374            token,375            run_id,376            chunk=None,377            parent_run_id=parent_run_id,378        )379 380    def _on_chat_model_start(self, run: Run) -> None:381        """Persist a chat model run.382 383        Note:384            Naming is historical: there is no `_on_chat_model_end` hook. Chat385            model completion is handled by `_on_llm_end`, shared with text386            LLM runs.387        """388        if run.parent_run_id is None:389            run.reference_example_id = self.example_id390        self._persist_run_single(run)391 392    def _on_llm_end(self, run: Run) -> None:393        """Process LLM/chat model run completion."""394        # Extract usage_metadata from outputs and store in extra.metadata395        if run.outputs and "generations" in run.outputs:396            usage_metadata = _get_usage_metadata_from_generations(397                run.outputs["generations"]398            )399            if usage_metadata is not None:400                if "metadata" not in run.extra:401                    run.extra["metadata"] = {}402                run.extra["metadata"]["usage_metadata"] = usage_metadata403        self._update_run_single(run)404 405    def _on_llm_error(self, run: Run) -> None:406        """Process the LLM Run upon error."""407        self._update_run_single(run)408 409    def _on_chain_start(self, run: Run) -> None:410        """Process the Chain Run upon start."""411        if run.parent_run_id is None:412            run.reference_example_id = self.example_id413        # Skip persisting if inputs are deferred (e.g., iterator/generator inputs).414        # The run will be posted when _on_chain_end is called with realized inputs.415        if not run.extra.get("defers_inputs"):416            self._persist_run_single(run)417 418    def _on_chain_end(self, run: Run) -> None:419        """Process the Chain Run."""420        # If inputs were deferred, persist (POST) the run now that inputs are realized.421        # Otherwise, update (PATCH) the existing run.422        if run.extra.get("defers_inputs"):423            self._persist_run_single(run)424        else:425            self._update_run_single(run)426 427    def _on_chain_error(self, run: Run) -> None:428        """Process the Chain Run upon error."""429        # If inputs were deferred, persist (POST) the run now that inputs are realized.430        # Otherwise, update (PATCH) the existing run.431        if run.extra.get("defers_inputs"):432            self._persist_run_single(run)433        else:434            self._update_run_single(run)435 436    def _on_tool_start(self, run: Run) -> None:437        """Process the Tool Run upon start."""438        if run.parent_run_id is None:439            run.reference_example_id = self.example_id440        self._persist_run_single(run)441 442    def _on_tool_end(self, run: Run) -> None:443        """Process the Tool Run."""444        self._update_run_single(run)445 446    def _on_tool_error(self, run: Run) -> None:447        """Process the Tool Run upon error."""448        self._update_run_single(run)449 450    def _on_retriever_start(self, run: Run) -> None:451        """Process the Retriever Run upon start."""452        if run.parent_run_id is None:453            run.reference_example_id = self.example_id454        self._persist_run_single(run)455 456    def _on_retriever_end(self, run: Run) -> None:457        """Process the Retriever Run."""458        self._update_run_single(run)459 460    def _on_retriever_error(self, run: Run) -> None:461        """Process the Retriever Run upon error."""462        self._update_run_single(run)463 464    def wait_for_futures(self) -> None:465        """Wait for the given futures to complete."""466        if self.client is not None:467            self.client.flush()468 469 470def _patch_missing_metadata(self: LangChainTracer, run: Run) -> None:471    if not self.tracing_metadata:472        return473    metadata = run.metadata474    patched = None475    for k, v in self.tracing_metadata.items():476        # ``OVERRIDABLE_LANGSMITH_INHERITABLE_METADATA_KEYS`` are a small,477        # LangSmith-only allowlist that bypasses the "first wins" merge478        # so a nested caller (e.g. a subagent) can override a parent-set value.479        if k not in metadata or k in OVERRIDABLE_LANGSMITH_INHERITABLE_METADATA_KEYS:480            # Skip the copy when the value already matches (avoids cloning481            # the shared dict in the common "already set" case). Use a482            # ``k in metadata`` guard so a legitimate missing key whose483            # tracer value happens to be ``None`` is still patched in.484            if k in metadata and metadata[k] == v:485                continue486            if patched is None:487                # Copy on first miss to avoid mutating the shared dict.488                patched = {**metadata}489                run.extra["metadata"] = patched490            patched[k] = v491