mirror of
https://github.com/hwchase17/langchain.git
synced 2025-05-19 14:01:50 +00:00
5fc44989bf
Thank you for contributing to LangChain!
- [ ] **PR title**: "package: description"
- Where "package" is whichever of langchain, community, core,
experimental, etc. is being modified. Use "docs: ..." for purely docs
changes, "templates: ..." for template changes, "infra: ..." for CI
changes.
- Example: "community: add foobar LLM"
- [ ] **PR message**: ***Delete this entire checklist*** and replace
with
- **Description:** a description of the change
- **Issue:** the issue # it fixes, if applicable
- **Dependencies:** any dependencies required for this change
- **Twitter handle:** if your PR gets announced, and you'd like a
mention, we'll gladly shout you out!
- [ ] **Add tests and docs**: If you're adding a new integration, please
include
1. a test for the integration, preferably unit tests that do not rely on
network access,
2. an example notebook showing its use. It lives in
`docs/docs/integrations` directory.
- [ ] **Lint and test**: Run `make format`, `make lint` and `make test`
from the root of the package(s) you've modified. See contribution
guidelines for more: https://python.langchain.com/docs/contributing/
Additional guidelines:
- Make sure optional dependencies are imported within a function.
- Please do not add dependencies to pyproject.toml files (even optional
ones) unless they are required for unit tests.
- Most PRs should not touch more than one package.
- Changes should be backwards compatible.
- If you are adding something to community, do not re-import it in
langchain.
If no one reviews your PR within a few days, please @-mention one of
baskaryan, efriis, eyurtsev, ccurme, vbarda, hwchase17.
---------
Co-authored-by: Erick Friis <erick@langchain.dev>
125 lines
5.5 KiB
Python
125 lines
5.5 KiB
Python
"""Test CallbackManager."""
|
|
|
|
from unittest.mock import patch
|
|
|
|
import pytest
|
|
from langchain_core.callbacks.manager import CallbackManager, trace_as_chain_group
|
|
from langchain_core.outputs import LLMResult
|
|
from langchain_core.tracers.langchain import LangChainTracer, wait_for_all_tracers
|
|
from langchain_core.utils.pydantic import get_fields
|
|
from langsmith import utils as ls_utils
|
|
|
|
from langchain_community.callbacks import get_openai_callback
|
|
from langchain_community.callbacks.manager import get_bedrock_anthropic_callback
|
|
from langchain_community.llms.openai import BaseOpenAI
|
|
|
|
|
|
def test_callback_manager_configure_context_vars(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""Test callback manager configuration."""
|
|
ls_utils.get_env_var.cache_clear()
|
|
ls_utils.get_tracer_project.cache_clear()
|
|
monkeypatch.setenv("LANGCHAIN_TRACING_V2", "true")
|
|
monkeypatch.setenv("LANGCHAIN_TRACING", "false")
|
|
monkeypatch.setenv("LANGCHAIN_API_KEY", "foo")
|
|
with patch.object(LangChainTracer, "_update_run_single"):
|
|
with patch.object(LangChainTracer, "_persist_run_single"):
|
|
with trace_as_chain_group("test") as group_manager:
|
|
assert len(group_manager.handlers) == 1
|
|
tracer = group_manager.handlers[0]
|
|
assert isinstance(tracer, LangChainTracer)
|
|
|
|
with get_openai_callback() as cb:
|
|
# This is a new empty callback handler
|
|
assert cb.successful_requests == 0
|
|
assert cb.total_tokens == 0
|
|
|
|
# configure adds this openai cb but doesn't modify the group manager
|
|
mngr = CallbackManager.configure(group_manager)
|
|
assert mngr.handlers == [tracer, cb]
|
|
assert group_manager.handlers == [tracer]
|
|
|
|
response = LLMResult(
|
|
generations=[],
|
|
llm_output={
|
|
"token_usage": {
|
|
"prompt_tokens": 2,
|
|
"completion_tokens": 1,
|
|
"total_tokens": 3,
|
|
},
|
|
"model_name": get_fields(BaseOpenAI)["model_name"].default,
|
|
},
|
|
)
|
|
mngr.on_llm_start({}, ["prompt"])[0].on_llm_end(response)
|
|
|
|
# The callback handler has been updated
|
|
assert cb.successful_requests == 1
|
|
assert cb.total_tokens == 3
|
|
assert cb.prompt_tokens == 2
|
|
assert cb.completion_tokens == 1
|
|
assert cb.total_cost > 0
|
|
|
|
with get_openai_callback() as cb:
|
|
# This is a new empty callback handler
|
|
assert cb.successful_requests == 0
|
|
assert cb.total_tokens == 0
|
|
|
|
# configure adds this openai cb but doesn't modify the group manager
|
|
mngr = CallbackManager.configure(group_manager)
|
|
assert mngr.handlers == [tracer, cb]
|
|
assert group_manager.handlers == [tracer]
|
|
|
|
response = LLMResult(
|
|
generations=[],
|
|
llm_output={
|
|
"token_usage": {
|
|
"prompt_tokens": 2,
|
|
"completion_tokens": 1,
|
|
"total_tokens": 3,
|
|
},
|
|
"model_name": get_fields(BaseOpenAI)["model_name"].default,
|
|
},
|
|
)
|
|
mngr.on_llm_start({}, ["prompt"])[0].on_llm_end(response)
|
|
|
|
# The callback handler has been updated
|
|
assert cb.successful_requests == 1
|
|
assert cb.total_tokens == 3
|
|
assert cb.prompt_tokens == 2
|
|
assert cb.completion_tokens == 1
|
|
assert cb.total_cost > 0
|
|
|
|
with get_bedrock_anthropic_callback() as cb:
|
|
# This is a new empty callback handler
|
|
assert cb.successful_requests == 0
|
|
assert cb.total_tokens == 0
|
|
|
|
# configure adds this bedrock anthropic cb,
|
|
# but doesn't modify the group manager
|
|
mngr = CallbackManager.configure(group_manager)
|
|
assert mngr.handlers == [tracer, cb]
|
|
assert group_manager.handlers == [tracer]
|
|
|
|
response = LLMResult(
|
|
generations=[],
|
|
llm_output={
|
|
"usage": {
|
|
"prompt_tokens": 2,
|
|
"completion_tokens": 1,
|
|
"total_tokens": 3,
|
|
},
|
|
"model_id": "anthropic.claude-instant-v1",
|
|
},
|
|
)
|
|
mngr.on_llm_start({}, ["prompt"])[0].on_llm_end(response)
|
|
|
|
# The callback handler has been updated
|
|
assert cb.successful_requests == 1
|
|
assert cb.total_tokens == 3
|
|
assert cb.prompt_tokens == 2
|
|
assert cb.completion_tokens == 1
|
|
assert cb.total_cost > 0
|
|
wait_for_all_tracers()
|
|
assert LangChainTracer._persist_run_single.call_count == 4 # type: ignore
|