From e9d7726a6f8c49dbec593437008282249adb7575 Mon Sep 17 00:00:00 2001 From: lievan <42917263+lievan@users.noreply.github.com> Date: Wed, 29 May 2024 11:30:18 -0400 Subject: [PATCH] fix(llmobs): fix langchain nested llm spans (#9354) - Fix nested LLM Spans for LLMObs langchain integration when openai/bedrock is enabled - If the LLMObs integration detects that one of our supported LLM integrations is enabled (bedrock or openai) and also detects that the user is using a Langchain plugin for these llm integrations (which we detect using the `model_provider` tag set by the langchain integration), then the llm operation is not marked as a workflow span. If not, the llm operation is marked as a an llm span. We also add a `TestLangchainTraceStructureWithLlmIntegrations` class to allow us to assert on the "trace structure" (where structure is defined by the order of span kinds enqueued to the LLM Obs trace writer) to verify the changes above ## Bedrock ### With bedrock enabled image ### Without bedrock enabled image ## OpenAI ### With openai enabled: image ### Without openai enabled image ## Checklist - [x] Change(s) are motivated and described in the PR description - [x] Testing strategy is described if automated tests are not included in the PR - [x] Risks are described (performance impact, potential for breakage, maintainability) - [x] Change is maintainable (easy to change, telemetry, documentation) - [x] [Library release note guidelines](https://ddtrace.readthedocs.io/en/stable/releasenotes.html) are followed or label `changelog/no-changelog` is set - [x] Documentation is included (in-code, generated user docs, [public corp docs](https://github.com/DataDog/documentation/)) - [x] Backport labels are set (if [applicable](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting)) - [x] If this PR changes the public interface, I've notified `@DataDog/apm-tees`. ## Reviewer Checklist - [ ] Title is accurate - [ ] All changes are related to the pull request's stated goal - [ ] Description motivates each change - [ ] Avoids breaking [API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces) changes - [ ] Testing strategy adequately addresses listed risks - [ ] Change is maintainable (easy to change, telemetry, documentation) - [ ] Release note makes sense to a user of the library - [ ] Author has acknowledged and discussed the performance implications of this PR as reported in the benchmarks PR comment - [ ] Backport labels are set in a manner that is consistent with the [release branch maintenance policy](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting) --------- Co-authored-by: lievan Co-authored-by: Brett Langdon --- .riot/requirements/12739ce.txt | 91 ++++++++++ .riot/requirements/15757fd.txt | 79 --------- .riot/requirements/1bd8794.txt | 81 --------- .riot/requirements/1e00937.txt | 81 --------- .riot/requirements/6251000.txt | 86 +++++++++ .riot/requirements/da7b0f6.txt | 88 ++++++++++ ddtrace/llmobs/_integrations/langchain.py | 42 +++-- ddtrace/llmobs/_llmobs.py | 6 + riotfile.py | 14 +- .../bedrock_amazon_chat_invoke.yaml | 73 ++++++++ .../bedrock_amazon_invoke.yaml | 58 ++++++ .../langchain/test_langchain_community.py | 166 ++++++++++++++++++ .../contrib/langchain/test_langchain_patch.py | 2 + 13 files changed, 610 insertions(+), 257 deletions(-) create mode 100644 .riot/requirements/12739ce.txt delete mode 100644 .riot/requirements/15757fd.txt delete mode 100644 .riot/requirements/1bd8794.txt delete mode 100644 .riot/requirements/1e00937.txt create mode 100644 .riot/requirements/6251000.txt create mode 100644 .riot/requirements/da7b0f6.txt create mode 100644 tests/contrib/langchain/cassettes/langchain_community/bedrock_amazon_chat_invoke.yaml create mode 100644 tests/contrib/langchain/cassettes/langchain_community/bedrock_amazon_invoke.yaml diff --git a/.riot/requirements/12739ce.txt b/.riot/requirements/12739ce.txt new file mode 100644 index 00000000000..d0ceae82d1d --- /dev/null +++ b/.riot/requirements/12739ce.txt @@ -0,0 +1,91 @@ +# +# This file is autogenerated by pip-compile with Python 3.9 +# by the following command: +# +# pip-compile --no-annotate .riot/requirements/12739ce.in +# +ai21==2.4.0 +ai21-tokenizer==0.9.1 +aiohttp==3.9.5 +aiosignal==1.3.1 +annotated-types==0.7.0 +anyio==4.4.0 +async-timeout==4.0.3 +attrs==23.2.0 +boto3==1.34.114 +botocore==1.34.114 +certifi==2024.2.2 +charset-normalizer==3.3.2 +cohere==5.5.3 +coverage[toml]==7.5.3 +dataclasses-json==0.6.6 +distro==1.9.0 +exceptiongroup==1.2.1 +fastavro==1.9.4 +filelock==3.14.0 +frozenlist==1.4.1 +fsspec==2024.5.0 +h11==0.14.0 +httpcore==1.0.5 +httpx==0.27.0 +httpx-sse==0.4.0 +huggingface-hub==0.23.2 +hypothesis==6.45.0 +idna==3.7 +importlib-metadata==7.1.0 +iniconfig==2.0.0 +jmespath==1.0.1 +jsonpatch==1.33 +jsonpointer==2.4 +langchain==0.1.20 +langchain-aws==0.1.6 +langchain-community==0.0.38 +langchain-core==0.1.52 +langchain-openai==0.1.5 +langchain-pinecone==0.1.1 +langchain-text-splitters==0.0.2 +langsmith==0.1.63 +marshmallow==3.21.2 +mock==5.1.0 +multidict==6.0.5 +mypy-extensions==1.0.0 +numexpr==2.10.0 +numpy==1.26.4 +openai==1.12.0 +opentracing==2.4.0 +orjson==3.10.3 +packaging==23.2 +pinecone-client==3.2.2 +pluggy==1.5.0 +psutil==5.9.8 +pydantic==2.7.2 +pydantic-core==2.18.3 +pytest==8.2.1 +pytest-asyncio==0.21.1 +pytest-cov==5.0.0 +pytest-mock==3.14.0 +pytest-randomly==3.15.0 +python-dateutil==2.9.0.post0 +pyyaml==6.0.1 +regex==2024.5.15 +requests==2.32.2 +s3transfer==0.10.1 +sentencepiece==0.2.0 +six==1.16.0 +sniffio==1.3.1 +sortedcontainers==2.4.0 +sqlalchemy==2.0.30 +tenacity==8.3.0 +tiktoken==0.7.0 +tokenizers==0.19.1 +tomli==2.0.1 +tqdm==4.66.4 +types-requests==2.31.0.6 +types-urllib3==1.26.25.14 +typing-extensions==4.12.0 +typing-inspect==0.9.0 +urllib3==1.26.18 +vcrpy==6.0.1 +wrapt==1.16.0 +yarl==1.9.4 +zipp==3.19.0 diff --git a/.riot/requirements/15757fd.txt b/.riot/requirements/15757fd.txt deleted file mode 100644 index f5293e2208d..00000000000 --- a/.riot/requirements/15757fd.txt +++ /dev/null @@ -1,79 +0,0 @@ -# -# This file is autogenerated by pip-compile with Python 3.11 -# by the following command: -# -# pip-compile --no-annotate .riot/requirements/15757fd.in -# -ai21==2.1.2 -ai21-tokenizer==0.3.11 -aiohttp==3.9.3 -aiosignal==1.3.1 -annotated-types==0.6.0 -anyio==4.3.0 -attrs==23.2.0 -backoff==2.2.1 -certifi==2024.2.2 -charset-normalizer==3.3.2 -cohere==4.53 -coverage[toml]==7.4.3 -dataclasses-json==0.6.4 -distro==1.9.0 -exceptiongroup==1.2.0 -fastavro==1.9.4 -filelock==3.13.1 -frozenlist==1.4.1 -fsspec==2024.2.0 -greenlet==3.0.3 -h11==0.14.0 -httpcore==1.0.4 -httpx==0.27.0 -huggingface-hub==0.21.4 -hypothesis==6.45.0 -idna==3.6 -importlib-metadata==6.11.0 -iniconfig==2.0.0 -jsonpatch==1.33 -jsonpointer==2.4 -langchain==0.1.9 -langchain-community==0.0.24 -langchain-core==0.1.27 -langchain-openai==0.0.8 -langchain-pinecone==0.0.3 -langsmith==0.1.9 -marshmallow==3.21.1 -mock==5.1.0 -multidict==6.0.5 -mypy-extensions==1.0.0 -numexpr==2.9.0 -numpy==1.26.4 -openai==1.12.0 -opentracing==2.4.0 -orjson==3.9.15 -packaging==23.2 -pinecone-client==3.1.0 -pluggy==1.4.0 -psutil==5.9.8 -pydantic==2.6.3 -pydantic-core==2.16.3 -pytest==8.1.1 -pytest-asyncio==0.21.1 -pytest-cov==4.1.0 -pytest-mock==3.12.0 -pytest-randomly==3.15.0 -pyyaml==6.0.1 -regex==2023.12.25 -requests==2.31.0 -sentencepiece==0.1.99 -sniffio==1.3.1 -sortedcontainers==2.4.0 -sqlalchemy==2.0.28 -tenacity==8.2.3 -tiktoken==0.6.0 -tqdm==4.66.2 -typing-extensions==4.10.0 -typing-inspect==0.9.0 -urllib3==2.2.1 -vcrpy==6.0.1 -wrapt==1.16.0 -yarl==1.9.4 -zipp==3.17.0 diff --git a/.riot/requirements/1bd8794.txt b/.riot/requirements/1bd8794.txt deleted file mode 100644 index 2b606dfceee..00000000000 --- a/.riot/requirements/1bd8794.txt +++ /dev/null @@ -1,81 +0,0 @@ -# -# This file is autogenerated by pip-compile with Python 3.10 -# by the following command: -# -# pip-compile --no-annotate .riot/requirements/1bd8794.in -# -ai21==2.1.2 -ai21-tokenizer==0.3.11 -aiohttp==3.9.3 -aiosignal==1.3.1 -annotated-types==0.6.0 -anyio==4.3.0 -async-timeout==4.0.3 -attrs==23.2.0 -backoff==2.2.1 -certifi==2024.2.2 -charset-normalizer==3.3.2 -cohere==4.53 -coverage[toml]==7.4.3 -dataclasses-json==0.6.4 -distro==1.9.0 -exceptiongroup==1.2.0 -fastavro==1.9.4 -filelock==3.13.1 -frozenlist==1.4.1 -fsspec==2024.2.0 -greenlet==3.0.3 -h11==0.14.0 -httpcore==1.0.4 -httpx==0.27.0 -huggingface-hub==0.21.4 -hypothesis==6.45.0 -idna==3.6 -importlib-metadata==6.11.0 -iniconfig==2.0.0 -jsonpatch==1.33 -jsonpointer==2.4 -langchain==0.1.9 -langchain-community==0.0.24 -langchain-core==0.1.27 -langchain-openai==0.0.8 -langchain-pinecone==0.0.3 -langsmith==0.1.9 -marshmallow==3.21.1 -mock==5.1.0 -multidict==6.0.5 -mypy-extensions==1.0.0 -numexpr==2.9.0 -numpy==1.26.4 -openai==1.12.0 -opentracing==2.4.0 -orjson==3.9.15 -packaging==23.2 -pinecone-client==3.1.0 -pluggy==1.4.0 -psutil==5.9.8 -pydantic==2.6.3 -pydantic-core==2.16.3 -pytest==8.1.1 -pytest-asyncio==0.21.1 -pytest-cov==4.1.0 -pytest-mock==3.12.0 -pytest-randomly==3.15.0 -pyyaml==6.0.1 -regex==2023.12.25 -requests==2.31.0 -sentencepiece==0.1.99 -sniffio==1.3.1 -sortedcontainers==2.4.0 -sqlalchemy==2.0.28 -tenacity==8.2.3 -tiktoken==0.6.0 -tomli==2.0.1 -tqdm==4.66.2 -typing-extensions==4.10.0 -typing-inspect==0.9.0 -urllib3==2.2.1 -vcrpy==6.0.1 -wrapt==1.16.0 -yarl==1.9.4 -zipp==3.17.0 diff --git a/.riot/requirements/1e00937.txt b/.riot/requirements/1e00937.txt deleted file mode 100644 index ed20d38b60f..00000000000 --- a/.riot/requirements/1e00937.txt +++ /dev/null @@ -1,81 +0,0 @@ -# -# This file is autogenerated by pip-compile with Python 3.9 -# by the following command: -# -# pip-compile --no-annotate .riot/requirements/1e00937.in -# -ai21==2.1.2 -ai21-tokenizer==0.3.11 -aiohttp==3.9.3 -aiosignal==1.3.1 -annotated-types==0.6.0 -anyio==4.3.0 -async-timeout==4.0.3 -attrs==23.2.0 -backoff==2.2.1 -certifi==2024.2.2 -charset-normalizer==3.3.2 -cohere==4.53 -coverage[toml]==7.4.3 -dataclasses-json==0.6.4 -distro==1.9.0 -exceptiongroup==1.2.0 -fastavro==1.9.4 -filelock==3.13.1 -frozenlist==1.4.1 -fsspec==2024.2.0 -greenlet==3.0.3 -h11==0.14.0 -httpcore==1.0.4 -httpx==0.27.0 -huggingface-hub==0.21.4 -hypothesis==6.45.0 -idna==3.6 -importlib-metadata==6.11.0 -iniconfig==2.0.0 -jsonpatch==1.33 -jsonpointer==2.4 -langchain==0.1.9 -langchain-community==0.0.24 -langchain-core==0.1.27 -langchain-openai==0.0.8 -langchain-pinecone==0.0.3 -langsmith==0.1.9 -marshmallow==3.21.1 -mock==5.1.0 -multidict==6.0.5 -mypy-extensions==1.0.0 -numexpr==2.9.0 -numpy==1.26.4 -openai==1.12.0 -opentracing==2.4.0 -orjson==3.9.15 -packaging==23.2 -pinecone-client==3.1.0 -pluggy==1.4.0 -psutil==5.9.8 -pydantic==2.6.3 -pydantic-core==2.16.3 -pytest==8.1.1 -pytest-asyncio==0.21.1 -pytest-cov==4.1.0 -pytest-mock==3.12.0 -pytest-randomly==3.15.0 -pyyaml==6.0.1 -regex==2023.12.25 -requests==2.31.0 -sentencepiece==0.1.99 -sniffio==1.3.1 -sortedcontainers==2.4.0 -sqlalchemy==2.0.28 -tenacity==8.2.3 -tiktoken==0.6.0 -tomli==2.0.1 -tqdm==4.66.2 -typing-extensions==4.10.0 -typing-inspect==0.9.0 -urllib3==1.26.18 -vcrpy==6.0.1 -wrapt==1.16.0 -yarl==1.9.4 -zipp==3.17.0 diff --git a/.riot/requirements/6251000.txt b/.riot/requirements/6251000.txt new file mode 100644 index 00000000000..78abfa4a9c2 --- /dev/null +++ b/.riot/requirements/6251000.txt @@ -0,0 +1,86 @@ +# +# This file is autogenerated by pip-compile with Python 3.11 +# by the following command: +# +# pip-compile --no-annotate .riot/requirements/6251000.in +# +ai21==2.4.0 +ai21-tokenizer==0.9.1 +aiohttp==3.9.5 +aiosignal==1.3.1 +annotated-types==0.7.0 +anyio==4.4.0 +attrs==23.2.0 +boto3==1.34.114 +botocore==1.34.114 +certifi==2024.2.2 +charset-normalizer==3.3.2 +cohere==5.5.3 +coverage[toml]==7.5.3 +dataclasses-json==0.6.6 +distro==1.9.0 +exceptiongroup==1.2.1 +fastavro==1.9.4 +filelock==3.14.0 +frozenlist==1.4.1 +fsspec==2024.5.0 +h11==0.14.0 +httpcore==1.0.5 +httpx==0.27.0 +httpx-sse==0.4.0 +huggingface-hub==0.23.2 +hypothesis==6.45.0 +idna==3.7 +iniconfig==2.0.0 +jmespath==1.0.1 +jsonpatch==1.33 +jsonpointer==2.4 +langchain==0.1.20 +langchain-aws==0.1.6 +langchain-community==0.0.38 +langchain-core==0.1.52 +langchain-openai==0.1.5 +langchain-pinecone==0.1.1 +langchain-text-splitters==0.0.2 +langsmith==0.1.63 +marshmallow==3.21.2 +mock==5.1.0 +multidict==6.0.5 +mypy-extensions==1.0.0 +numexpr==2.10.0 +numpy==1.26.4 +openai==1.12.0 +opentracing==2.4.0 +orjson==3.10.3 +packaging==23.2 +pinecone-client==3.2.2 +pluggy==1.5.0 +psutil==5.9.8 +pydantic==2.7.2 +pydantic-core==2.18.3 +pytest==8.2.1 +pytest-asyncio==0.21.1 +pytest-cov==5.0.0 +pytest-mock==3.14.0 +pytest-randomly==3.15.0 +python-dateutil==2.9.0.post0 +pyyaml==6.0.1 +regex==2024.5.15 +requests==2.32.2 +s3transfer==0.10.1 +sentencepiece==0.2.0 +six==1.16.0 +sniffio==1.3.1 +sortedcontainers==2.4.0 +sqlalchemy==2.0.30 +tenacity==8.3.0 +tiktoken==0.7.0 +tokenizers==0.19.1 +tqdm==4.66.4 +types-requests==2.32.0.20240523 +typing-extensions==4.12.0 +typing-inspect==0.9.0 +urllib3==2.2.1 +vcrpy==6.0.1 +wrapt==1.16.0 +yarl==1.9.4 diff --git a/.riot/requirements/da7b0f6.txt b/.riot/requirements/da7b0f6.txt new file mode 100644 index 00000000000..f2c28511799 --- /dev/null +++ b/.riot/requirements/da7b0f6.txt @@ -0,0 +1,88 @@ +# +# This file is autogenerated by pip-compile with Python 3.10 +# by the following command: +# +# pip-compile --no-annotate .riot/requirements/da7b0f6.in +# +ai21==2.4.0 +ai21-tokenizer==0.9.1 +aiohttp==3.9.5 +aiosignal==1.3.1 +annotated-types==0.7.0 +anyio==4.4.0 +async-timeout==4.0.3 +attrs==23.2.0 +boto3==1.34.114 +botocore==1.34.114 +certifi==2024.2.2 +charset-normalizer==3.3.2 +cohere==5.5.3 +coverage[toml]==7.5.3 +dataclasses-json==0.6.6 +distro==1.9.0 +exceptiongroup==1.2.1 +fastavro==1.9.4 +filelock==3.14.0 +frozenlist==1.4.1 +fsspec==2024.5.0 +h11==0.14.0 +httpcore==1.0.5 +httpx==0.27.0 +httpx-sse==0.4.0 +huggingface-hub==0.23.2 +hypothesis==6.45.0 +idna==3.7 +iniconfig==2.0.0 +jmespath==1.0.1 +jsonpatch==1.33 +jsonpointer==2.4 +langchain==0.1.20 +langchain-aws==0.1.6 +langchain-community==0.0.38 +langchain-core==0.1.52 +langchain-openai==0.1.5 +langchain-pinecone==0.1.1 +langchain-text-splitters==0.0.2 +langsmith==0.1.63 +marshmallow==3.21.2 +mock==5.1.0 +multidict==6.0.5 +mypy-extensions==1.0.0 +numexpr==2.10.0 +numpy==1.26.4 +openai==1.12.0 +opentracing==2.4.0 +orjson==3.10.3 +packaging==23.2 +pinecone-client==3.2.2 +pluggy==1.5.0 +psutil==5.9.8 +pydantic==2.7.2 +pydantic-core==2.18.3 +pytest==8.2.1 +pytest-asyncio==0.21.1 +pytest-cov==5.0.0 +pytest-mock==3.14.0 +pytest-randomly==3.15.0 +python-dateutil==2.9.0.post0 +pyyaml==6.0.1 +regex==2024.5.15 +requests==2.32.2 +s3transfer==0.10.1 +sentencepiece==0.2.0 +six==1.16.0 +sniffio==1.3.1 +sortedcontainers==2.4.0 +sqlalchemy==2.0.30 +tenacity==8.3.0 +tiktoken==0.7.0 +tokenizers==0.19.1 +tomli==2.0.1 +tqdm==4.66.4 +types-requests==2.32.0.20240523 +typing-extensions==4.12.0 +typing-inspect==0.9.0 +urllib3==2.2.1 +vcrpy==6.0.1 +wrapt==1.16.0 +yarl==1.9.4 diff --git a/ddtrace/llmobs/_integrations/langchain.py b/ddtrace/llmobs/_integrations/langchain.py index d6a8d7e9ea7..c55e5d6085b 100644 --- a/ddtrace/llmobs/_integrations/langchain.py +++ b/ddtrace/llmobs/_integrations/langchain.py @@ -9,6 +9,7 @@ from ddtrace._trace.span import Span from ddtrace.constants import ERROR_TYPE from ddtrace.internal.logger import get_logger +from ddtrace.llmobs import LLMObs from ddtrace.llmobs._constants import INPUT_MESSAGES from ddtrace.llmobs._constants import INPUT_VALUE from ddtrace.llmobs._constants import METADATA @@ -31,6 +32,9 @@ TOTAL_COST = "langchain.tokens.total_cost" TYPE = "langchain.request.type" +BEDROCK_PROVIDER_NAME = "amazon_bedrock" +OPENAI_PROVIDER_NAME = "openai" + ROLE_MAPPING = { "human": "user", "ai": "assistant", @@ -55,13 +59,23 @@ def llmobs_set_tags( model_provider = span.get_tag(PROVIDER) self._llmobs_set_metadata(span, model_provider) + is_workflow = False + + if model_provider: + llmobs_integration = "custom" + if model_provider.startswith(BEDROCK_PROVIDER_NAME): + llmobs_integration = "bedrock" + elif model_provider.startswith(OPENAI_PROVIDER_NAME): + llmobs_integration = "openai" + + is_workflow = LLMObs._integration_is_enabled(llmobs_integration) + if operation == "llm": - self._llmobs_set_meta_tags_from_llm(span, inputs, response, error) + self._llmobs_set_meta_tags_from_llm(span, inputs, response, error, is_workflow=is_workflow) elif operation == "chat": - self._llmobs_set_meta_tags_from_chat_model(span, inputs, response, error) + self._llmobs_set_meta_tags_from_chat_model(span, inputs, response, error, is_workflow=is_workflow) elif operation == "chain": self._llmobs_set_meta_tags_from_chain(span, inputs, response, error) - span.set_tag_str(METRICS, json.dumps({})) def _llmobs_set_metadata(self, span: Span, model_provider: Optional[str] = None) -> None: @@ -86,20 +100,24 @@ def _llmobs_set_metadata(self, span: Span, model_provider: Optional[str] = None) span.set_tag_str(METADATA, json.dumps(metadata)) def _llmobs_set_meta_tags_from_llm( - self, span: Span, prompts: List[Any], completions: Any, err: bool = False + self, span: Span, prompts: List[Any], completions: Any, err: bool = False, is_workflow: bool = False ) -> None: - span.set_tag_str(SPAN_KIND, "llm") + span.set_tag_str(SPAN_KIND, "workflow" if is_workflow else "llm") span.set_tag_str(MODEL_NAME, span.get_tag(MODEL) or "") span.set_tag_str(MODEL_PROVIDER, span.get_tag(PROVIDER) or "") + input_tag_key = INPUT_VALUE if is_workflow else INPUT_MESSAGES + output_tag_key = OUTPUT_VALUE if is_workflow else OUTPUT_MESSAGES + if isinstance(prompts, str): prompts = [prompts] - span.set_tag_str(INPUT_MESSAGES, json.dumps([{"content": str(prompt)} for prompt in prompts])) + + span.set_tag_str(input_tag_key, json.dumps([{"content": str(prompt)} for prompt in prompts])) message_content = [{"content": ""}] if not err: message_content = [{"content": completion[0].text} for completion in completions.generations] - span.set_tag_str(OUTPUT_MESSAGES, json.dumps(message_content)) + span.set_tag_str(output_tag_key, json.dumps(message_content)) def _llmobs_set_meta_tags_from_chat_model( self, @@ -107,11 +125,15 @@ def _llmobs_set_meta_tags_from_chat_model( chat_messages: List[List[Any]], chat_completions: Any, err: bool = False, + is_workflow: bool = False, ) -> None: - span.set_tag_str(SPAN_KIND, "llm") + span.set_tag_str(SPAN_KIND, "workflow" if is_workflow else "llm") span.set_tag_str(MODEL_NAME, span.get_tag(MODEL) or "") span.set_tag_str(MODEL_PROVIDER, span.get_tag(PROVIDER) or "") + input_tag_key = INPUT_VALUE if is_workflow else INPUT_MESSAGES + output_tag_key = OUTPUT_VALUE if is_workflow else OUTPUT_MESSAGES + input_messages = [] for message_set in chat_messages: for message in message_set: @@ -122,7 +144,7 @@ def _llmobs_set_meta_tags_from_chat_model( "role": getattr(message, "role", ROLE_MAPPING.get(message.type, "")), } ) - span.set_tag_str(INPUT_MESSAGES, json.dumps(input_messages)) + span.set_tag_str(input_tag_key, json.dumps(input_messages)) output_messages = [{"content": ""}] if not err: @@ -137,7 +159,7 @@ def _llmobs_set_meta_tags_from_chat_model( "role": role, } ) - span.set_tag_str(OUTPUT_MESSAGES, json.dumps(output_messages)) + span.set_tag_str(output_tag_key, json.dumps(output_messages)) def _llmobs_set_meta_tags_from_chain( self, diff --git a/ddtrace/llmobs/_llmobs.py b/ddtrace/llmobs/_llmobs.py index 0476e91803b..26c6c5db4f1 100644 --- a/ddtrace/llmobs/_llmobs.py +++ b/ddtrace/llmobs/_llmobs.py @@ -174,6 +174,12 @@ def enable( atexit.register(cls.disable) log.debug("%s enabled", cls.__name__) + @classmethod + def _integration_is_enabled(cls, integration): + if integration not in SUPPORTED_LLMOBS_INTEGRATIONS: + return False + return SUPPORTED_LLMOBS_INTEGRATIONS[integration] in ddtrace._monkey._get_patched_modules() + @classmethod def disable(cls) -> None: if not cls.enabled: diff --git a/riotfile.py b/riotfile.py index dcdffe51162..17a34dffac4 100644 --- a/riotfile.py +++ b/riotfile.py @@ -2503,14 +2503,16 @@ def select_pys(min_version=MIN_PYTHON_VERSION, max_version=MAX_PYTHON_VERSION): ), Venv( pkgs={ - "langchain": "==0.1.9", - "langchain-community": "==0.0.24", - "langchain-core": "==0.1.27", - "langchain-openai": "==0.0.8", - "langchain-pinecone": "==0.0.3", - "langsmith": "==0.1.9", + "langchain": latest, + "langchain-community": latest, + "langchain-core": latest, + "langchain-openai": latest, + "langchain-pinecone": latest, + "langsmith": latest, "openai": "==1.12.0", "pinecone-client": latest, + "botocore": latest, + "langchain-aws": latest, } ), ], diff --git a/tests/contrib/langchain/cassettes/langchain_community/bedrock_amazon_chat_invoke.yaml b/tests/contrib/langchain/cassettes/langchain_community/bedrock_amazon_chat_invoke.yaml new file mode 100644 index 00000000000..2791bb834cd --- /dev/null +++ b/tests/contrib/langchain/cassettes/langchain_community/bedrock_amazon_chat_invoke.yaml @@ -0,0 +1,73 @@ +interactions: +- request: + body: '{"inputText": "\n\nUser: summarize the plot to the lord of the rings in + a dozen words\n\nBot:", "textGenerationConfig": {}}' + headers: + Accept: + - !!binary | + YXBwbGljYXRpb24vanNvbg== + Content-Length: + - '123' + Content-Type: + - !!binary | + YXBwbGljYXRpb24vanNvbg== + User-Agent: + - !!binary | + Qm90bzMvMS4zNC4xMTIgbWQvQm90b2NvcmUjMS4zNC4xMTIgdWEvMi4wIG9zL21hY29zIzIzLjMu + MCBtZC9hcmNoI2FybTY0IGxhbmcvcHl0aG9uIzMuMTAuMTMgbWQvcHlpbXBsI0NQeXRob24gY2Zn + L3JldHJ5LW1vZGUjbGVnYWN5IEJvdG9jb3JlLzEuMzQuMTEy + X-Amz-Date: + - !!binary | + MjAyNDA1MjRUMjA0ODI3Wg== + X-Amz-Security-Token: + - !!binary | + SVFvSmIzSnBaMmx1WDJWakVCMGFDWFZ6TFdWaGMzUXRNU0pJTUVZQ0lRRGllZ0VUWjBnamZCaFZS + b3A5RVkvRnRyc2s0ZjJCNzhyajUzRWZCQ2s3QVFJaEFLbThHM2h4OEtMb3cwV0d1QW1TcmhoeDJi + OUtWTEgzZ2pCaXVqT0V3NEZXS3BFRENKYi8vLy8vLy8vLy93RVFBeG9NTmpBeE5ESTNNamM1T1Rr + d0lneU5uYitTOTBWaWw4Undlc2txNVFJYm1SMEczcldHbWZiUVRmbGpOQmRlSlQxMzFZT2JqNEQw + YStkRHE4WWdoLzBMVEtFQkoyVGpoczlESTBhRkVXeDJlTnh1U0xkTkxYT2VaYWNDRDR3WDh4UjEy + ZzFRMlprVm9TQ0szS0N6UnAyTVMrcFlWR2h5WU4xOC8xZkZnOWlsR1FDTlA1SHFRc0E2R2g4U2U0 + SEdZTGl4dE5HNVJjRGxnTHI4aklMaG1PM3dtOGY2dVhKWkRmdjhHcGNDRk50YlNJNmFiVTkrdlFV + dytlT3l0NVoza0pQZDZYbklsZ3pvZ0tQd3Y0Wk9yZTMycThvdDl0aldXSjk1VDdzZkI2UTZsdkl6 + M1h2eFNqMDl0Vmhic3pubVE3aDZiZHp2T2YyYTFkRWx6RU40ZnVRaVhrZXNZTjQrdVYxVjdvNXVN + ODFKRXVGOHBoQWtRMS9OYmV1ZTJIVkRieXh2dUpuU0ZnY0hQN2pJUDQyOXFqcWIzOVl5R3I5QUFl + UHRtRDV2amhFOUt4emxrRTdvQTFCVWtaRE9CNWVIN0YxanJUZDcvcmdoWVJwOE1MVFhIVElsdHd1 + RnJxRjYwY0UwRFZoUHphbHFWMzZvZlhvU1lMWmsyUnNURVF2amtOQUVMM0F3NVBqRHNnWTZwUUdx + YkYrOWtvK250T2RHck1zNUltSEZOME1SU2ZmcmwxMEs0N2hhcGtWUDFoQXVCRlBrZEpzc0FySzhZ + RDlqQWN2TzNjT0M5b3dORDdBdWc2RERRaUFKMGVveldGbU9BbXZtbTNIbWlFS2swc3BscDNMeXJV + UUVtWjUwUmI5SmJNRndLc3pRaGgyYml2b3NhTm02bE55M005NlIzVi9odU1EMVpCZ2FWamJxSjBD + WDJobFlidDdsR0lsdHFUNi9jSVdibWxDV0wzUG5EUGs3MHJEakM2Vi9IZ3VkeTBFPQ== + amz-sdk-invocation-id: + - !!binary | + NjJmOTVjMWItMTg4ZS00MmQzLTgzOGMtOTNiNmJjNDAyYmE5 + amz-sdk-request: + - !!binary | + YXR0ZW1wdD0x + method: POST + uri: https://bedrock-runtime.us-east-1.amazonaws.com/model/amazon.titan-tg1-large/invoke + response: + body: + string: '{"inputTextTokenCount":21,"results":[{"tokenCount":28,"outputText":" + The plot of The Lord of the Rings is about a hobbit named Frodo Baggins who + is tasked with destroying a magical ring.","completionReason":"FINISH"}]}' + headers: + Connection: + - keep-alive + Content-Length: + - '218' + Content-Type: + - application/json + Date: + - Fri, 24 May 2024 20:48:29 GMT + X-Amzn-Bedrock-Input-Token-Count: + - '21' + X-Amzn-Bedrock-Invocation-Latency: + - '1478' + X-Amzn-Bedrock-Output-Token-Count: + - '28' + x-amzn-RequestId: + - 805deac9-65cd-4564-b98f-8e77001aa48f + status: + code: 200 + message: OK +version: 1 diff --git a/tests/contrib/langchain/cassettes/langchain_community/bedrock_amazon_invoke.yaml b/tests/contrib/langchain/cassettes/langchain_community/bedrock_amazon_invoke.yaml new file mode 100644 index 00000000000..ba7fddc88ed --- /dev/null +++ b/tests/contrib/langchain/cassettes/langchain_community/bedrock_amazon_invoke.yaml @@ -0,0 +1,58 @@ +interactions: +- request: + body: '{"inputText": "Command: can you explain what Datadog is to someone not + in the tech industry?", "textGenerationConfig": {"maxTokenCount": 50, "stopSequences": + [], "temperature": 0, "topP": 0.9}}' + headers: + Accept: + - !!binary | + YXBwbGljYXRpb24vanNvbg== + Content-Length: + - '193' + Content-Type: + - !!binary | + YXBwbGljYXRpb24vanNvbg== + User-Agent: + - !!binary | + Qm90bzMvMS4zNC42IG1kL0JvdG9jb3JlIzEuMzQuNiB1YS8yLjAgb3MvbWFjb3MjMjMuMi4wIG1k + L2FyY2gjYXJtNjQgbGFuZy9weXRob24jMy4xMC41IG1kL3B5aW1wbCNDUHl0aG9uIGNmZy9yZXRy + eS1tb2RlI2xlZ2FjeSBCb3RvY29yZS8xLjM0LjY= + X-Amz-Date: + - !!binary | + MjAyNDAxMDhUMTgyNTIzWg== + amz-sdk-invocation-id: + - !!binary | + NTNjODllZTEtYWUwZS00MWRlLWFiZGUtMmQ1NGYxZDQ1MWEw + amz-sdk-request: + - !!binary | + YXR0ZW1wdD0x + method: POST + uri: https://bedrock-runtime.us-east-1.amazonaws.com/model/amazon.titan-tg1-large/invoke + response: + body: + string: '{"inputTextTokenCount":18,"results":[{"tokenCount":50,"outputText":"\nDatadog + is a monitoring and analytics platform that helps businesses track and optimize + their digital infrastructure. It provides a comprehensive set of tools and + services for monitoring server performance, application health, and user behavior. + With Datadog, businesses can identify performance","completionReason":"LENGTH"}]}' + headers: + Connection: + - keep-alive + Content-Length: + - '397' + Content-Type: + - application/json + Date: + - Mon, 08 Jan 2024 18:25:26 GMT + X-Amzn-Bedrock-Input-Token-Count: + - '18' + X-Amzn-Bedrock-Invocation-Latency: + - '1835' + X-Amzn-Bedrock-Output-Token-Count: + - '50' + x-amzn-RequestId: + - 758fa023-a298-48d0-89e6-0208306cb76d + status: + code: 200 + message: OK +version: 1 diff --git a/tests/contrib/langchain/test_langchain_community.py b/tests/contrib/langchain/test_langchain_community.py index 88f4223db7c..cf946f9d981 100644 --- a/tests/contrib/langchain/test_langchain_community.py +++ b/tests/contrib/langchain/test_langchain_community.py @@ -14,6 +14,8 @@ from tests.contrib.langchain.utils import get_request_vcr from tests.llmobs._utils import _expected_llmobs_llm_span_event from tests.llmobs._utils import _expected_llmobs_non_llm_span_event +from tests.subprocesstest import SubprocessTestCase +from tests.subprocesstest import run_in_subprocess from tests.utils import flaky from tests.utils import override_global_config @@ -1621,3 +1623,167 @@ def test_llmobs_chain_schema_io( ), ], ) + + +class TestLangchainTraceStructureWithLlmIntegrations(SubprocessTestCase): + bedrock_env_config = dict( + AWS_ACCESS_KEY_ID="testing", + AWS_SECRET_ACCESS_KEY="testing", + AWS_SECURITY_TOKEN="testing", + AWS_SESSION_TOKEN="testing", + AWS_DEFAULT_REGION="us-east-1", + DD_LANGCHAIN_METRICS_ENABLED="false", + DD_API_KEY="", + ) + + openai_env_config = dict( + OPENAI_API_KEY="testing", + DD_API_KEY="", + ) + + def setUp(self): + patcher = mock.patch("ddtrace.llmobs._llmobs.LLMObsSpanWriter") + LLMObsSpanWriterMock = patcher.start() + mock_llmobs_span_writer = mock.MagicMock() + LLMObsSpanWriterMock.return_value = mock_llmobs_span_writer + + self.mock_llmobs_span_writer = mock_llmobs_span_writer + + super(TestLangchainTraceStructureWithLlmIntegrations, self).setUp() + + def _assert_trace_structure_from_writer_call_args(self, span_kinds): + assert self.mock_llmobs_span_writer.enqueue.call_count == len(span_kinds) + + calls = self.mock_llmobs_span_writer.enqueue.call_args_list + + for span_kind, call in zip(span_kinds, calls): + call_args = call.args[0] + + assert call_args["meta"]["span.kind"] == span_kind + if span_kind == "workflow": + assert len(call_args["meta"]["input"]["value"]) > 0 + assert len(call_args["meta"]["output"]["value"]) > 0 + elif span_kind == "llm": + assert len(call_args["meta"]["input"]["messages"]) > 0 + assert len(call_args["meta"]["output"]["messages"]) > 0 + + def _call_bedrock_chat_model(self, ChatBedrock, HumanMessage): + chat = ChatBedrock( + model_id="amazon.titan-tg1-large", + model_kwargs={"max_tokens": 50, "temperature": 0}, + ) + messages = [HumanMessage(content="summarize the plot to the lord of the rings in a dozen words")] + with get_request_vcr(subdirectory_name="langchain_community").use_cassette("bedrock_amazon_chat_invoke.yaml"): + chat.invoke(messages) + + def _call_bedrock_llm(self, Bedrock, ConversationChain, ConversationBufferMemory): + llm = Bedrock( + model_id="amazon.titan-tg1-large", + region_name="us-east-1", + model_kwargs={"temperature": 0, "topP": 0.9, "stopSequences": [], "maxTokens": 50}, + ) + + conversation = ConversationChain(llm=llm, verbose=True, memory=ConversationBufferMemory()) + + with get_request_vcr(subdirectory_name="langchain_community").use_cassette("bedrock_amazon_invoke.yaml"): + conversation.predict(input="can you explain what Datadog is to someone not in the tech industry?") + + def _call_openai_llm(self, OpenAI): + llm = OpenAI() + with get_request_vcr(subdirectory_name="langchain_community").use_cassette("openai_completion_sync.yaml"): + llm.invoke("Can you explain what Descartes meant by 'I think, therefore I am'?") + + @run_in_subprocess(env_overrides=bedrock_env_config) + def test_llmobs_with_chat_model_bedrock_enabled(self): + from langchain_aws import ChatBedrock + from langchain_core.messages import HumanMessage + + from ddtrace import patch + from ddtrace.llmobs import LLMObs + + patch(langchain=True, botocore=True) + LLMObs.enable(ml_app="", integrations_enabled=False, agentless_enabled=True) + + self._call_bedrock_chat_model(ChatBedrock, HumanMessage) + + self._assert_trace_structure_from_writer_call_args(["workflow", "llm"]) + + LLMObs.disable() + + @run_in_subprocess(env_overrides=bedrock_env_config) + def test_llmobs_with_chat_model_bedrock_disabled(self): + from langchain_aws import ChatBedrock + from langchain_core.messages import HumanMessage + + from ddtrace import patch + from ddtrace.llmobs import LLMObs + + patch(langchain=True) + LLMObs.enable(ml_app="", integrations_enabled=False, agentless_enabled=True) + + self._call_bedrock_chat_model(ChatBedrock, HumanMessage) + + self._assert_trace_structure_from_writer_call_args(["llm"]) + + LLMObs.disable() + + @run_in_subprocess(env_overrides=bedrock_env_config) + def test_llmobs_with_llm_model_bedrock_enabled(self): + from langchain.chains import ConversationChain + from langchain.memory import ConversationBufferMemory + from langchain_community.llms import Bedrock + + from ddtrace import patch + from ddtrace.llmobs import LLMObs + + patch(langchain=True, botocore=True) + LLMObs.enable(ml_app="", integrations_enabled=False, agentless_enabled=True) + self._call_bedrock_llm(Bedrock, ConversationChain, ConversationBufferMemory) + self._assert_trace_structure_from_writer_call_args(["workflow", "workflow", "llm"]) + + LLMObs.disable() + + @run_in_subprocess(env_overrides=bedrock_env_config) + def test_llmobs_with_llm_model_bedrock_disabled(self): + from langchain.chains import ConversationChain + from langchain.memory import ConversationBufferMemory + from langchain_community.llms import Bedrock + + from ddtrace import patch + from ddtrace.llmobs import LLMObs + + patch(langchain=True) + LLMObs.enable(ml_app="", integrations_enabled=False, agentless_enabled=True) + self._call_bedrock_llm(Bedrock, ConversationChain, ConversationBufferMemory) + self._assert_trace_structure_from_writer_call_args(["workflow", "llm"]) + + LLMObs.disable() + + @run_in_subprocess(env_overrides=openai_env_config) + def test_llmobs_langchain_with_openai_enabled(self): + from langchain_openai import OpenAI + + from ddtrace import patch + from ddtrace.llmobs import LLMObs + + patch(langchain=True, openai=True) + LLMObs.enable(ml_app="", integrations_enabled=False, agentless_enabled=True) + self._call_openai_llm(OpenAI) + self._assert_trace_structure_from_writer_call_args(["workflow", "llm"]) + + LLMObs.disable() + + @run_in_subprocess(env_overrides=openai_env_config) + def test_llmobs_langchain_with_openai_disabled(self): + from langchain_openai import OpenAI + + from ddtrace import patch + from ddtrace.llmobs import LLMObs + + patch(langchain=True) + + LLMObs.enable(ml_app="", integrations_enabled=False, agentless_enabled=True) + self._call_openai_llm(OpenAI) + self._assert_trace_structure_from_writer_call_args(["llm"]) + + LLMObs.disable() diff --git a/tests/contrib/langchain/test_langchain_patch.py b/tests/contrib/langchain/test_langchain_patch.py index 5d937d2c28a..dd17e6e7781 100644 --- a/tests/contrib/langchain/test_langchain_patch.py +++ b/tests/contrib/langchain/test_langchain_patch.py @@ -55,6 +55,7 @@ def assert_module_patched(self, langchain): def assert_not_module_patched(self, langchain): if SHOULD_PATCH_LANGCHAIN_COMMUNITY: from langchain import chains # noqa: F401 + from langchain.chains import base # noqa: F401 import langchain_community as gated_langchain from langchain_community import embeddings # noqa: F401 from langchain_community import vectorstores # noqa: F401 @@ -95,6 +96,7 @@ def assert_not_module_patched(self, langchain): def assert_not_module_double_patched(self, langchain): if SHOULD_PATCH_LANGCHAIN_COMMUNITY: + from langchain.chains import base # noqa: F401 import langchain_community as gated_langchain import langchain_core import langchain_openai