test_async_base_tracer.py — langchain Source File

Source Code

"""Test Tracer classes."""

from __future__ import annotations

from datetime import datetime, timezone
from typing import Any
from uuid import uuid4

import pytest
from freezegun import freeze_time

from langchain_core.callbacks import AsyncCallbackManager
from langchain_core.exceptions import TracerException
from langchain_core.messages import HumanMessage
from langchain_core.outputs import LLMResult
from langchain_core.tracers._compat import pydantic_to_dict
from langchain_core.tracers.base import AsyncBaseTracer
from langchain_core.tracers.schemas import Run

SERIALIZED = {"id": ["llm"]}
SERIALIZED_CHAT = {"id": ["chat_model"]}


class FakeAsyncTracer(AsyncBaseTracer):
    """Fake tracer to test async based tracers."""

    def __init__(self) -> None:
        """Initialize the tracer."""
        super().__init__()
        self.runs: list[Run] = []

    async def _persist_run(self, run: Run) -> None:
        self.runs.append(run)


def _compare_run_with_error(run: Any, expected_run: Any) -> None:
    if run.child_runs:
        assert len(expected_run.child_runs) == len(run.child_runs)
        for received, expected in zip(
            run.child_runs, expected_run.child_runs, strict=False
        ):
            _compare_run_with_error(received, expected)
    received = pydantic_to_dict(run, exclude={"child_runs"})
    received_err = received.pop("error")
    expected = pydantic_to_dict(expected_run, exclude={"child_runs"})
    expected_err = expected.pop("error")

    assert received == expected
    if expected_err is not None:
        assert received_err is not None
        assert expected_err in received_err
    else:
        assert received_err is None


@freeze_time("2023-01-01")
async def test_tracer_llm_run() -> None:
    """Test tracer on an LLM run."""
    uuid = uuid4()
    compare_run = Run(
// ... (556 more lines)