test_on_llm_end_preserves_existing_metadata() — langchain Function Reference

Source Code

libs/core/tests/unit_tests/tracers/test_langchain.py lines 316–355

def test_on_llm_end_preserves_existing_metadata() -> None:
    """Test that existing metadata is preserved when adding `usage_metadata`."""
    client = unittest.mock.MagicMock(spec=Client)
    client.tracing_queue = None
    tracer = LangChainTracer(client=client)

    run_id = UUID("9d878ab3-e5ca-4218-aef6-44cbdc90160a")
    tracer.on_llm_start(
        {"name": "test_llm"},
        ["foo"],
        run_id=run_id,
        metadata={"existing_key": "existing_value"},
    )

    run = tracer.run_map[str(run_id)]
    usage_metadata = {"input_tokens": 10, "output_tokens": 20, "total_tokens": 30}
    run.outputs = {
        "generations": [
            [
                {
                    "text": "Hello!",
                    "message": {"content": "Hello!", "usage_metadata": usage_metadata},
                }
            ]
        ]
    }

    captured_run = None

    def capture_run(r: Run) -> None:
        nonlocal captured_run
        captured_run = r

    with unittest.mock.patch.object(tracer, "_update_run_single", capture_run):
        tracer._on_llm_end(run)

    assert captured_run is not None
    assert "metadata" in captured_run.extra
    assert captured_run.extra["metadata"]["usage_metadata"] == usage_metadata
    assert captured_run.extra["metadata"]["existing_key"] == "existing_value"