codekingpro/portable-devtools

1"""Base interfaces for tracing runs."""2 3from __future__ import annotations4 5import asyncio6import logging7from abc import ABC, abstractmethod8from typing import (9    TYPE_CHECKING,10    Any,11)12 13from typing_extensions import override14 15from langchain_core.callbacks.base import AsyncCallbackHandler, BaseCallbackHandler16from langchain_core.exceptions import TracerException  # noqa: F40117from langchain_core.tracers.core import _TracerCore18 19if TYPE_CHECKING:20    from collections.abc import Sequence21    from uuid import UUID22 23    from tenacity import RetryCallState24 25    from langchain_core.documents import Document26    from langchain_core.messages import BaseMessage27    from langchain_core.outputs import ChatGenerationChunk, GenerationChunk, LLMResult28    from langchain_core.tracers.schemas import Run29 30logger = logging.getLogger(__name__)31 32 33class BaseTracer(_TracerCore, BaseCallbackHandler, ABC):34    """Base interface for tracers."""35 36    @abstractmethod37    def _persist_run(self, run: Run) -> None:38        """Persist a run."""39 40    def _start_trace(self, run: Run) -> None:41        """Start a trace for a run."""42        super()._start_trace(run)43        self._on_run_create(run)44 45    def _end_trace(self, run: Run) -> None:46        """End a trace for a run."""47        if not run.parent_run_id:48            self._persist_run(run)49        self.run_map.pop(str(run.id))50        # If this run's parent was injected from an external tracing context51        # (e.g. a langsmith @traceable), decrement its child refcount and52        # remove it from run_map once the last child is done.53        parent_id = str(run.parent_run_id) if run.parent_run_id else None54        if parent_id and parent_id in self._external_run_ids:55            self._external_run_ids[parent_id] -= 156            if self._external_run_ids[parent_id] <= 0:57                self.run_map.pop(parent_id, None)58                del self._external_run_ids[parent_id]59        self._on_run_update(run)60 61    def on_chat_model_start(62        self,63        serialized: dict[str, Any],64        messages: list[list[BaseMessage]],65        *,66        run_id: UUID,67        tags: list[str] | None = None,68        parent_run_id: UUID | None = None,69        metadata: dict[str, Any] | None = None,70        name: str | None = None,71        **kwargs: Any,72    ) -> Run:73        """Start a trace for a chat model run.74 75        Note:76            Naming can be confusing here: there is `on_chat_model_start`, but no77            corresponding `on_chat_model_end` callback. Chat model completion is78            routed through `on_llm_end` / `_on_llm_end`, which are shared with79            text LLM runs.80 81        Args:82            serialized: The serialized model.83            messages: The messages to start the chat with.84            run_id: The run ID.85            tags: The tags for the run.86            parent_run_id: The parent run ID.87            metadata: The metadata for the run.88            name: The name of the run.89            **kwargs: Additional arguments.90 91        Returns:92            The run.93        """94        chat_model_run = self._create_chat_model_run(95            serialized=serialized,96            messages=messages,97            run_id=run_id,98            parent_run_id=parent_run_id,99            tags=tags,100            metadata=metadata,101            name=name,102            **kwargs,103        )104        self._start_trace(chat_model_run)105        self._on_chat_model_start(chat_model_run)106        return chat_model_run107 108    def on_llm_start(109        self,110        serialized: dict[str, Any],111        prompts: list[str],112        *,113        run_id: UUID,114        tags: list[str] | None = None,115        parent_run_id: UUID | None = None,116        metadata: dict[str, Any] | None = None,117        name: str | None = None,118        **kwargs: Any,119    ) -> Run:120        """Start a trace for an LLM run.121 122        Args:123            serialized: The serialized model.124            prompts: The prompts to start the LLM with.125            run_id: The run ID.126            tags: The tags for the run.127            parent_run_id: The parent run ID.128            metadata: The metadata for the run.129            name: The name of the run.130            **kwargs: Additional arguments.131 132        Returns:133            The run.134        """135        llm_run = self._create_llm_run(136            serialized=serialized,137            prompts=prompts,138            run_id=run_id,139            parent_run_id=parent_run_id,140            tags=tags,141            metadata=metadata,142            name=name,143            **kwargs,144        )145        self._start_trace(llm_run)146        self._on_llm_start(llm_run)147        return llm_run148 149    @override150    def on_llm_new_token(151        self,152        token: str,153        *,154        chunk: GenerationChunk | ChatGenerationChunk | None = None,155        run_id: UUID,156        parent_run_id: UUID | None = None,157        **kwargs: Any,158    ) -> Run:159        """Run on new LLM token.160 161        Only available when streaming is enabled.162 163        Args:164            token: The token.165            chunk: The chunk.166            run_id: The run ID.167            parent_run_id: The parent run ID.168            **kwargs: Additional arguments.169 170        Returns:171            The run.172        """173        # "chat_model" is only used for the experimental new streaming_events format.174        # This change should not affect any existing tracers.175        llm_run = self._llm_run_with_token_event(176            token=token,177            run_id=run_id,178            chunk=chunk,179            parent_run_id=parent_run_id,180        )181        self._on_llm_new_token(llm_run, token, chunk)182        return llm_run183 184    @override185    def on_retry(186        self,187        retry_state: RetryCallState,188        *,189        run_id: UUID,190        **kwargs: Any,191    ) -> Run:192        """Run on retry.193 194        Args:195            retry_state: The retry state.196            run_id: The run ID.197            **kwargs: Additional arguments.198 199        Returns:200            The run.201        """202        return self._llm_run_with_retry_event(203            retry_state=retry_state,204            run_id=run_id,205        )206 207    @override208    def on_llm_end(self, response: LLMResult, *, run_id: UUID, **kwargs: Any) -> Run:209        """End a trace for an LLM or chat model run.210 211        Note:212            This is the end callback for both run types. Chat models start with213            `on_chat_model_start`, but there is no `on_chat_model_end`;214            completion is routed here for callback API compatibility.215 216        Args:217            response: The response.218            run_id: The run ID.219            **kwargs: Additional arguments.220 221        Returns:222            The run.223        """224        # "chat_model" is only used for the experimental new streaming_events format.225        # This change should not affect any existing tracers.226        llm_run = self._complete_llm_run(227            response=response,228            run_id=run_id,229        )230        self._end_trace(llm_run)231        self._on_llm_end(llm_run)232        return llm_run233 234    def on_llm_error(235        self,236        error: BaseException,237        *,238        run_id: UUID,239        **kwargs: Any,240    ) -> Run:241        """Handle an error for an LLM run.242 243        Args:244            error: The error.245            run_id: The run ID.246            **kwargs: Additional arguments.247 248        Returns:249            The run.250        """251        # "chat_model" is only used for the experimental new streaming_events format.252        # This change should not affect any existing tracers.253        llm_run = self._errored_llm_run(254            error=error, run_id=run_id, response=kwargs.pop("response", None)255        )256        self._end_trace(llm_run)257        self._on_llm_error(llm_run)258        return llm_run259 260    @override261    def on_chain_start(262        self,263        serialized: dict[str, Any],264        inputs: dict[str, Any],265        *,266        run_id: UUID,267        tags: list[str] | None = None,268        parent_run_id: UUID | None = None,269        metadata: dict[str, Any] | None = None,270        run_type: str | None = None,271        name: str | None = None,272        **kwargs: Any,273    ) -> Run:274        """Start a trace for a chain run.275 276        Args:277            serialized: The serialized chain.278            inputs: The inputs for the chain.279            run_id: The run ID.280            tags: The tags for the run.281            parent_run_id: The parent run ID.282            metadata: The metadata for the run.283            run_type: The type of the run.284            name: The name of the run.285            **kwargs: Additional arguments.286 287        Returns:288            The run.289        """290        chain_run = self._create_chain_run(291            serialized=serialized,292            inputs=inputs,293            run_id=run_id,294            tags=tags,295            parent_run_id=parent_run_id,296            metadata=metadata,297            run_type=run_type,298            name=name,299            **kwargs,300        )301        self._start_trace(chain_run)302        self._on_chain_start(chain_run)303        return chain_run304 305    @override306    def on_chain_end(307        self,308        outputs: dict[str, Any],309        *,310        run_id: UUID,311        inputs: dict[str, Any] | None = None,312        **kwargs: Any,313    ) -> Run:314        """End a trace for a chain run.315 316        Args:317            outputs: The outputs for the chain.318            run_id: The run ID.319            inputs: The inputs for the chain.320            **kwargs: Additional arguments.321 322        Returns:323            The run.324        """325        chain_run = self._complete_chain_run(326            outputs=outputs,327            run_id=run_id,328            inputs=inputs,329        )330        self._end_trace(chain_run)331        self._on_chain_end(chain_run)332        return chain_run333 334    @override335    def on_chain_error(336        self,337        error: BaseException,338        *,339        inputs: dict[str, Any] | None = None,340        run_id: UUID,341        **kwargs: Any,342    ) -> Run:343        """Handle an error for a chain run.344 345        Args:346            error: The error.347            inputs: The inputs for the chain.348            run_id: The run ID.349            **kwargs: Additional arguments.350 351        Returns:352            The run.353        """354        chain_run = self._errored_chain_run(355            error=error,356            run_id=run_id,357            inputs=inputs,358        )359        self._end_trace(chain_run)360        self._on_chain_error(chain_run)361        return chain_run362 363    def on_tool_start(364        self,365        serialized: dict[str, Any],366        input_str: str,367        *,368        run_id: UUID,369        tags: list[str] | None = None,370        parent_run_id: UUID | None = None,371        metadata: dict[str, Any] | None = None,372        name: str | None = None,373        inputs: dict[str, Any] | None = None,374        **kwargs: Any,375    ) -> Run:376        """Start a trace for a tool run.377 378        Args:379            serialized: The serialized tool.380            input_str: The input string.381            run_id: The run ID.382            tags: The tags for the run.383            parent_run_id: The parent run ID.384            metadata: The metadata for the run.385            name: The name of the run.386            inputs: The inputs for the tool.387            **kwargs: Additional arguments.388 389        Returns:390            The run.391        """392        tool_run = self._create_tool_run(393            serialized=serialized,394            input_str=input_str,395            run_id=run_id,396            tags=tags,397            parent_run_id=parent_run_id,398            metadata=metadata,399            name=name,400            inputs=inputs,401            **kwargs,402        )403        self._start_trace(tool_run)404        self._on_tool_start(tool_run)405        return tool_run406 407    @override408    def on_tool_end(self, output: Any, *, run_id: UUID, **kwargs: Any) -> Run:409        """End a trace for a tool run.410 411        Args:412            output: The output for the tool.413            run_id: The run ID.414            **kwargs: Additional arguments.415 416        Returns:417            The run.418        """419        tool_run = self._complete_tool_run(420            output=output,421            run_id=run_id,422        )423        self._end_trace(tool_run)424        self._on_tool_end(tool_run)425        return tool_run426 427    @override428    def on_tool_error(429        self,430        error: BaseException,431        *,432        run_id: UUID,433        **kwargs: Any,434    ) -> Run:435        """Handle an error for a tool run.436 437        Args:438            error: The error.439            run_id: The run ID.440            **kwargs: Additional arguments.441 442        Returns:443            The run.444        """445        tool_run = self._errored_tool_run(446            error=error,447            run_id=run_id,448        )449        self._end_trace(tool_run)450        self._on_tool_error(tool_run)451        return tool_run452 453    def on_retriever_start(454        self,455        serialized: dict[str, Any],456        query: str,457        *,458        run_id: UUID,459        parent_run_id: UUID | None = None,460        tags: list[str] | None = None,461        metadata: dict[str, Any] | None = None,462        name: str | None = None,463        **kwargs: Any,464    ) -> Run:465        """Run when the `Retriever` starts running.466 467        Args:468            serialized: The serialized retriever.469            query: The query.470            run_id: The run ID.471            parent_run_id: The parent run ID.472            tags: The tags for the run.473            metadata: The metadata for the run.474            name: The name of the run.475            **kwargs: Additional arguments.476 477        Returns:478            The run.479        """480        retrieval_run = self._create_retrieval_run(481            serialized=serialized,482            query=query,483            run_id=run_id,484            parent_run_id=parent_run_id,485            tags=tags,486            metadata=metadata,487            name=name,488            **kwargs,489        )490        self._start_trace(retrieval_run)491        self._on_retriever_start(retrieval_run)492        return retrieval_run493 494    @override495    def on_retriever_error(496        self,497        error: BaseException,498        *,499        run_id: UUID,500        **kwargs: Any,501    ) -> Run:502        """Run when `Retriever` errors.503 504        Args:505            error: The error.506            run_id: The run ID.507            **kwargs: Additional arguments.508 509        Returns:510            The run.511        """512        retrieval_run = self._errored_retrieval_run(513            error=error,514            run_id=run_id,515        )516        self._end_trace(retrieval_run)517        self._on_retriever_error(retrieval_run)518        return retrieval_run519 520    @override521    def on_retriever_end(522        self, documents: Sequence[Document], *, run_id: UUID, **kwargs: Any523    ) -> Run:524        """Run when the `Retriever` ends running.525 526        Args:527            documents: The documents.528            run_id: The run ID.529            **kwargs: Additional arguments.530 531        Returns:532            The run.533        """534        retrieval_run = self._complete_retrieval_run(535            documents=documents,536            run_id=run_id,537        )538        self._end_trace(retrieval_run)539        self._on_retriever_end(retrieval_run)540        return retrieval_run541 542    def __deepcopy__(self, memo: dict) -> BaseTracer:543        """Return self."""544        return self545 546    def __copy__(self) -> BaseTracer:547        """Return self."""548        return self549 550 551class AsyncBaseTracer(_TracerCore, AsyncCallbackHandler, ABC):552    """Async base interface for tracers."""553 554    @abstractmethod555    @override556    async def _persist_run(self, run: Run) -> None:557        """Persist a run."""558 559    @override560    async def _start_trace(self, run: Run) -> None:561        """Start a trace for a run.562 563        Starting a trace will run concurrently with each `_on_[run_type]_start` method.564        No `_on_[run_type]_start` callback should depend on operations in565        `_start_trace`.566        """567        super()._start_trace(run)568        await self._on_run_create(run)569 570    @override571    async def _end_trace(self, run: Run) -> None:572        """End a trace for a run.573 574        Ending a trace will run concurrently with each `_on_[run_type]_end` method.575        No `_on_[run_type]_end` callback should depend on operations in `_end_trace`.576        """577        if not run.parent_run_id:578            await self._persist_run(run)579        self.run_map.pop(str(run.id))580        # If this run's parent was injected from an external tracing context581        # (e.g. a langsmith @traceable), decrement its child refcount and582        # remove it from run_map once the last child is done.583        parent_id = str(run.parent_run_id) if run.parent_run_id else None584        if parent_id and parent_id in self._external_run_ids:585            self._external_run_ids[parent_id] -= 1586            if self._external_run_ids[parent_id] <= 0:587                self.run_map.pop(parent_id, None)588                del self._external_run_ids[parent_id]589        await self._on_run_update(run)590 591    @override592    async def on_chat_model_start(593        self,594        serialized: dict[str, Any],595        messages: list[list[BaseMessage]],596        *,597        run_id: UUID,598        parent_run_id: UUID | None = None,599        tags: list[str] | None = None,600        metadata: dict[str, Any] | None = None,601        name: str | None = None,602        **kwargs: Any,603    ) -> Any:604        chat_model_run = self._create_chat_model_run(605            serialized=serialized,606            messages=messages,607            run_id=run_id,608            parent_run_id=parent_run_id,609            tags=tags,610            metadata=metadata,611            name=name,612            **kwargs,613        )614        tasks = [615            self._start_trace(chat_model_run),616            self._on_chat_model_start(chat_model_run),617        ]618        await asyncio.gather(*tasks)619        return chat_model_run620 621    @override622    async def on_llm_start(623        self,624        serialized: dict[str, Any],625        prompts: list[str],626        *,627        run_id: UUID,628        parent_run_id: UUID | None = None,629        tags: list[str] | None = None,630        metadata: dict[str, Any] | None = None,631        **kwargs: Any,632    ) -> None:633        llm_run = self._create_llm_run(634            serialized=serialized,635            prompts=prompts,636            run_id=run_id,637            parent_run_id=parent_run_id,638            tags=tags,639            metadata=metadata,640            **kwargs,641        )642        tasks = [self._start_trace(llm_run), self._on_llm_start(llm_run)]643        await asyncio.gather(*tasks)644 645    @override646    async def on_llm_new_token(647        self,648        token: str,649        *,650        chunk: GenerationChunk | ChatGenerationChunk | None = None,651        run_id: UUID,652        parent_run_id: UUID | None = None,653        **kwargs: Any,654    ) -> None:655        llm_run = self._llm_run_with_token_event(656            token=token,657            run_id=run_id,658            chunk=chunk,659            parent_run_id=parent_run_id,660        )661        await self._on_llm_new_token(llm_run, token, chunk)662 663    @override664    async def on_retry(665        self,666        retry_state: RetryCallState,667        *,668        run_id: UUID,669        **kwargs: Any,670    ) -> None:671        self._llm_run_with_retry_event(672            retry_state=retry_state,673            run_id=run_id,674        )675 676    @override677    async def on_llm_end(678        self,679        response: LLMResult,680        *,681        run_id: UUID,682        parent_run_id: UUID | None = None,683        tags: list[str] | None = None,684        **kwargs: Any,685    ) -> None:686        """End a trace for an LLM or chat model run.687 688        Note:689            This async callback also handles both run types. Async chat models690            start with `on_chat_model_start`, but there is no691            `on_chat_model_end`; completion is routed here for callback API692            compatibility.693        """694        llm_run = self._complete_llm_run(695            response=response,696            run_id=run_id,697        )698        tasks = [self._on_llm_end(llm_run), self._end_trace(llm_run)]699        await asyncio.gather(*tasks)700 701    @override702    async def on_llm_error(703        self,704        error: BaseException,705        *,706        run_id: UUID,707        parent_run_id: UUID | None = None,708        tags: list[str] | None = None,709        **kwargs: Any,710    ) -> None:711        llm_run = self._errored_llm_run(712            error=error,713            run_id=run_id,714        )715        tasks = [self._on_llm_error(llm_run), self._end_trace(llm_run)]716        await asyncio.gather(*tasks)717 718    @override719    async def on_chain_start(720        self,721        serialized: dict[str, Any],722        inputs: dict[str, Any],723        *,724        run_id: UUID,725        tags: list[str] | None = None,726        parent_run_id: UUID | None = None,727        metadata: dict[str, Any] | None = None,728        run_type: str | None = None,729        name: str | None = None,730        **kwargs: Any,731    ) -> None:732        chain_run = self._create_chain_run(733            serialized=serialized,734            inputs=inputs,735            run_id=run_id,736            tags=tags,737            parent_run_id=parent_run_id,738            metadata=metadata,739            run_type=run_type,740            name=name,741            **kwargs,742        )743        tasks = [self._start_trace(chain_run), self._on_chain_start(chain_run)]744        await asyncio.gather(*tasks)745 746    @override747    async def on_chain_end(748        self,749        outputs: dict[str, Any],750        *,751        run_id: UUID,752        inputs: dict[str, Any] | None = None,753        **kwargs: Any,754    ) -> None:755        chain_run = self._complete_chain_run(756            outputs=outputs,757            run_id=run_id,758            inputs=inputs,759        )760        tasks = [self._end_trace(chain_run), self._on_chain_end(chain_run)]761        await asyncio.gather(*tasks)762 763    @override764    async def on_chain_error(765        self,766        error: BaseException,767        *,768        inputs: dict[str, Any] | None = None,769        run_id: UUID,770        **kwargs: Any,771    ) -> None:772        chain_run = self._errored_chain_run(773            error=error,774            inputs=inputs,775            run_id=run_id,776        )777        tasks = [self._end_trace(chain_run), self._on_chain_error(chain_run)]778        await asyncio.gather(*tasks)779 780    @override781    async def on_tool_start(782        self,783        serialized: dict[str, Any],784        input_str: str,785        *,786        run_id: UUID,787        tags: list[str] | None = None,788        parent_run_id: UUID | None = None,789        metadata: dict[str, Any] | None = None,790        name: str | None = None,791        inputs: dict[str, Any] | None = None,792        **kwargs: Any,793    ) -> None:794        tool_run = self._create_tool_run(795            serialized=serialized,796            input_str=input_str,797            run_id=run_id,798            tags=tags,799            parent_run_id=parent_run_id,800            metadata=metadata,801            inputs=inputs,802            **kwargs,803        )804        tasks = [self._start_trace(tool_run), self._on_tool_start(tool_run)]805        await asyncio.gather(*tasks)806 807    @override808    async def on_tool_end(809        self,810        output: Any,811        *,812        run_id: UUID,813        **kwargs: Any,814    ) -> None:815        tool_run = self._complete_tool_run(816            output=output,817            run_id=run_id,818        )819        tasks = [self._end_trace(tool_run), self._on_tool_end(tool_run)]820        await asyncio.gather(*tasks)821 822    @override823    async def on_tool_error(824        self,825        error: BaseException,826        *,827        run_id: UUID,828        parent_run_id: UUID | None = None,829        tags: list[str] | None = None,830        **kwargs: Any,831    ) -> None:832        tool_run = self._errored_tool_run(833            error=error,834            run_id=run_id,835        )836        tasks = [self._end_trace(tool_run), self._on_tool_error(tool_run)]837        await asyncio.gather(*tasks)838 839    @override840    async def on_retriever_start(841        self,842        serialized: dict[str, Any],843        query: str,844        *,845        run_id: UUID,846        parent_run_id: UUID | None = None,847        tags: list[str] | None = None,848        metadata: dict[str, Any] | None = None,849        name: str | None = None,850        **kwargs: Any,851    ) -> None:852        retriever_run = self._create_retrieval_run(853            serialized=serialized,854            query=query,855            run_id=run_id,856            parent_run_id=parent_run_id,857            tags=tags,858            metadata=metadata,859            name=name,860        )861        tasks = [862            self._start_trace(retriever_run),863            self._on_retriever_start(retriever_run),864        ]865        await asyncio.gather(*tasks)866 867    @override868    async def on_retriever_error(869        self,870        error: BaseException,871        *,872        run_id: UUID,873        parent_run_id: UUID | None = None,874        tags: list[str] | None = None,875        **kwargs: Any,876    ) -> None:877        retrieval_run = self._errored_retrieval_run(878            error=error,879            run_id=run_id,880        )881        tasks = [882            self._end_trace(retrieval_run),883            self._on_retriever_error(retrieval_run),884        ]885        await asyncio.gather(*tasks)886 887    @override888    async def on_retriever_end(889        self,890        documents: Sequence[Document],891        *,892        run_id: UUID,893        parent_run_id: UUID | None = None,894        tags: list[str] | None = None,895        **kwargs: Any,896    ) -> None:897        retrieval_run = self._complete_retrieval_run(898            documents=documents,899            run_id=run_id,900        )901        tasks = [self._end_trace(retrieval_run), self._on_retriever_end(retrieval_run)]902        await asyncio.gather(*tasks)903 904    async def _on_run_create(self, run: Run) -> None:905        """Process a run upon creation."""906 907    async def _on_run_update(self, run: Run) -> None:908        """Process a run upon update."""909 910    async def _on_llm_start(self, run: Run) -> None:911        """Process the LLM Run upon start."""912 913    async def _on_llm_end(self, run: Run) -> None:914        """Process LLM/chat model run completion."""915 916    async def _on_llm_error(self, run: Run) -> None:917        """Process the LLM Run upon error."""918 919    async def _on_llm_new_token(920        self,921        run: Run,922        token: str,923        chunk: GenerationChunk | ChatGenerationChunk | None,924    ) -> None:925        """Process new LLM token."""926 927    async def _on_chain_start(self, run: Run) -> None:928        """Process the Chain Run upon start."""929 930    async def _on_chain_end(self, run: Run) -> None:931        """Process the Chain Run."""932 933    async def _on_chain_error(self, run: Run) -> None:934        """Process the Chain Run upon error."""935 936    async def _on_tool_start(self, run: Run) -> None:937        """Process the Tool Run upon start."""938 939    async def _on_tool_end(self, run: Run) -> None:940        """Process the Tool Run."""941 942    async def _on_tool_error(self, run: Run) -> None:943        """Process the Tool Run upon error."""944 945    async def _on_chat_model_start(self, run: Run) -> None:946        """Process the Chat Model Run upon start."""947 948    async def _on_retriever_start(self, run: Run) -> None:949        """Process the Retriever Run upon start."""950 951    async def _on_retriever_end(self, run: Run) -> None:952        """Process the Retriever Run."""953 954    async def _on_retriever_error(self, run: Run) -> None:955        """Process the Retriever Run upon error."""956