test_on_llm_end_stores_usage_metadata_in_run_extra() — langchain Function Reference

Source Code

libs/core/tests/unit_tests/tracers/test_langchain.py lines 252–285

def test_on_llm_end_stores_usage_metadata_in_run_extra() -> None:
    """Test that `usage_metadata` is stored in `run.extra.metadata` on llm end."""
    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)

    run = tracer.run_map[str(run_id)]
    usage_metadata = {"input_tokens": 100, "output_tokens": 200, "total_tokens": 300}
    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