From 4cfddeaf11841fc6e3c3fa2e66a3aef0837517db Mon Sep 17 00:00:00 2001 From: Sam Brenner Date: Mon, 15 Apr 2024 16:25:37 -0400 Subject: [PATCH] try test changes --- tests/contrib/langchain/test_langchain.py | 119 ++++++++++++++---- .../langchain/test_langchain_community.py | 107 +++++++++++++--- tests/llmobs/_utils.py | 2 +- 3 files changed, 182 insertions(+), 46 deletions(-) diff --git a/tests/contrib/langchain/test_langchain.py b/tests/contrib/langchain/test_langchain.py index 598711139f5..1b87d13f061 100644 --- a/tests/contrib/langchain/test_langchain.py +++ b/tests/contrib/langchain/test_langchain.py @@ -1266,7 +1266,7 @@ def test_vectorstore_logs_error(langchain, ddtrace_config_langchain, mock_logs, ) class TestLLMObsLangchain: @staticmethod - def _expected_llmobs_calls(trace, expected_spans_data: list): + def _expected_llmobs_chain_calls(trace, expected_spans_data: list): expected_llmobs_writer_calls = [mock.call.start()] for idx, span in enumerate(trace): @@ -1282,11 +1282,13 @@ def _expected_llmobs_calls(trace, expected_spans_data: list): return expected_llmobs_writer_calls @staticmethod - def _expected_llmobs_chain_call(span, input_parameters=None): + def _expected_llmobs_chain_call(span, input_parameters=None, input_value=None, output_value=None): return _expected_llmobs_non_llm_span_event( span, span_kind="workflow", parameters=input_parameters, + input_value=input_value, + output_value=output_value, tags={ "ml_app": "langchain_test", }, @@ -1333,7 +1335,44 @@ def _expected_llmobs_llm_calls(span, provider="openai", input_role=None, output_ ) @classmethod - def _test_llmobs_invoke( + def _test_llmobs_llm_invoke( + cls, + provider, + generate_trace, + request_vcr, + mock_llmobs_writer, + mock_tracer, + cassette_name, + input_role=None, + output_role=None, + different_py39_cassette=False, + ): + LLMObs.disable() + LLMObs.enable(tracer=mock_tracer) + + if sys.version_info < (3, 10, 0) and different_py39_cassette: + cassette_name = cassette_name.replace(".yaml", "_39.yaml") + with request_vcr.use_cassette(cassette_name): + generate_trace("Can you explain what an LLM chain is?") + span = mock_tracer.pop_traces()[0][0] + + expected_llmons_writer_calls = [ + mock.call.start(), + mock.call.enqueue( + cls._expected_llmobs_llm_calls( + span, + provider=provider, + input_role=input_role, + output_role=output_role, + ) + ), + ] + + assert mock_llmobs_writer.enqueue.call_count == 1 + mock_llmobs_writer.assert_has_calls(expected_llmons_writer_calls) + + @classmethod + def _test_llmobs_chain_invoke( cls, generate_trace, request_vcr, @@ -1352,44 +1391,47 @@ def _test_llmobs_invoke( generate_trace("Can you explain what an LLM chain is?") trace = mock_tracer.pop_traces()[0] - expected_llmobs_writer_calls = cls._expected_llmobs_calls(trace=trace, expected_spans_data=expected_spans_data) + expected_llmobs_writer_calls = cls._expected_llmobs_chain_calls( + trace=trace, expected_spans_data=expected_spans_data + ) assert mock_llmobs_writer.enqueue.call_count == len(expected_spans_data) mock_llmobs_writer.assert_has_calls(expected_llmobs_writer_calls) def test_llmobs_openai_llm(self, langchain, mock_llmobs_writer, mock_tracer, request_vcr): llm = langchain.llms.OpenAI() - self._test_llmobs_invoke( + self._test_llmobs_llm_invoke( generate_trace=llm, request_vcr=request_vcr, mock_llmobs_writer=mock_llmobs_writer, mock_tracer=mock_tracer, cassette_name="openai_completion_sync.yaml", different_py39_cassette=True, + provider="openai", ) def test_llmobs_cohere_llm(self, langchain, mock_llmobs_writer, mock_tracer, request_vcr): llm = langchain.llms.Cohere(model="cohere.command-light-text-v14") - self._test_llmobs_invoke( + self._test_llmobs_llm_invoke( generate_trace=llm, request_vcr=request_vcr, mock_llmobs_writer=mock_llmobs_writer, mock_tracer=mock_tracer, cassette_name="cohere_completion_sync.yaml", - expected_spans_data=[("llm", {"provider": "cohere", "input_role": None, "output_role": None})], + provider="cohere", ) def test_llmobs_ai21_llm(self, langchain, mock_llmobs_writer, mock_tracer, request_vcr): llm = langchain.llms.AI21() - self._test_llmobs_invoke( + self._test_llmobs_llm_invoke( generate_trace=llm, request_vcr=request_vcr, mock_llmobs_writer=mock_llmobs_writer, mock_tracer=mock_tracer, cassette_name="ai21_completion_sync.yaml", - expected_spans_data=[("llm", {"provider": "ai21", "input_role": None, "output_role": None})], + provider="ai21", different_py39_cassette=True, ) @@ -1400,59 +1442,67 @@ def test_llmobs_huggingfacehub_llm(self, langchain, mock_llmobs_writer, mock_tra huggingfacehub_api_token=os.getenv("HUGGINGFACEHUB_API_TOKEN", ""), ) - self._test_llmobs_invoke( + self._test_llmobs_llm_invoke( generate_trace=llm, request_vcr=request_vcr, mock_llmobs_writer=mock_llmobs_writer, mock_tracer=mock_tracer, cassette_name="huggingfacehub_completion_sync.yaml", - expected_spans_data=[("llm", {"provider": "huggingface_hub", "input_role": None, "output_role": None})], + provider="huggingface_hub", ) def test_llmobs_openai_chat_model(self, langchain, mock_llmobs_writer, mock_tracer, request_vcr): chat = langchain.chat_models.ChatOpenAI(temperature=0, max_tokens=256) - self._test_llmobs_invoke( + self._test_llmobs_llm_invoke( generate_trace=lambda prompt: chat([langchain.schema.HumanMessage(content=prompt)]), request_vcr=request_vcr, mock_llmobs_writer=mock_llmobs_writer, mock_tracer=mock_tracer, cassette_name="openai_chat_completion_sync_call.yaml", - expected_spans_data=[("llm", {"provider": "openai", "input_role": "user", "output_role": "assistant"})], + provider="openai", + input_role="user", + output_role="assistant", different_py39_cassette=True, ) def test_llmobs_openai_chat_model_custom_role(self, langchain, mock_llmobs_writer, mock_tracer, request_vcr): chat = langchain.chat_models.ChatOpenAI(temperature=0, max_tokens=256) - self._test_llmobs_invoke( + self._test_llmobs_llm_invoke( generate_trace=lambda prompt: chat([langchain.schema.ChatMessage(content=prompt, role="custom")]), request_vcr=request_vcr, mock_llmobs_writer=mock_llmobs_writer, mock_tracer=mock_tracer, cassette_name="openai_chat_completion_sync_call.yaml", - expected_spans_data=[("llm", {"provider": "openai", "input_role": "custom", "output_role": "assistant"})], + provider="openai", + input_role="custom", + output_role="assistant", different_py39_cassette=True, ) def test_llmobs_chain(self, langchain, mock_llmobs_writer, mock_tracer, request_vcr): chain = langchain.chains.LLMMathChain(llm=langchain.llms.OpenAI(temperature=0, max_tokens=256)) - self._test_llmobs_invoke( + self._test_llmobs_chain_invoke( generate_trace=lambda prompt: chain.run("what is two raised to the fifty-fourth power?"), request_vcr=request_vcr, mock_llmobs_writer=mock_llmobs_writer, mock_tracer=mock_tracer, cassette_name="openai_math_chain_sync.yaml", expected_spans_data=[ - ("chain", {"input_parameters": {"question": "what is two raised to the fifty-fourth power?"}}), ( "chain", { - "input_parameters": { - "question": "what is two raised to the fifty-fourth power?", - "stop": mock.ANY, - } + "input_value": str({"question": "what is two raised to the fifty-fourth power?"}), + "output_value": mock.ANY, + }, + ), + ( + "chain", + { + "input_value": mock.ANY, + "output_value": mock.ANY, }, ), ("llm", {"provider": "openai", "input_role": None, "output_role": None}), @@ -1495,17 +1545,36 @@ def test_llmobs_chain_nested(self, langchain, mock_llmobs_writer, mock_tracer, r true whenever it is put forward by me or conceived in my mind. """ - self._test_llmobs_invoke( + self._test_llmobs_chain_invoke( generate_trace=lambda prompt: sequential_chain.run({"input_text": input_text}), request_vcr=request_vcr, mock_llmobs_writer=mock_llmobs_writer, mock_tracer=mock_tracer, cassette_name="openai_sequential_paraphrase_and_rhyme_sync.yaml", expected_spans_data=[ - ("chain", {"input_parameters": {"input_text": input_text}}), - ("chain", {"input_parameters": {"input_text": input_text}}), + ( + "chain", + { + "input_value": str({"input_text": input_text}), + "output_value": mock.ANY, + }, + ), + ( + "chain", + { + "input_value": str({"input_text": input_text}), + "output_value": mock.ANY, + }, + ), ("llm", {"provider": "openai", "input_role": None, "output_role": None}), - ("chain", {"input_parameters": {"input_text": input_text, "paraphrased_output": mock.ANY}}), + ( + "chain", + { + "input_value": mock.ANY, + "output_value": mock.ANY, + }, + ), ("llm", {"provider": "openai", "input_role": None, "output_role": None}), ], + different_py39_cassette=True, ) diff --git a/tests/contrib/langchain/test_langchain_community.py b/tests/contrib/langchain/test_langchain_community.py index 76fac09a745..0c648dc3f96 100644 --- a/tests/contrib/langchain/test_langchain_community.py +++ b/tests/contrib/langchain/test_langchain_community.py @@ -1250,7 +1250,7 @@ async def test_lcel_chain_batch_async(langchain_core, langchain_openai, request_ ) class TestLLMObsLangchain: @staticmethod - def _expected_llmobs_calls(trace, expected_spans_data: list): + def _expected_llmobs_chain_calls(trace, expected_spans_data: list): expected_llmobs_writer_calls = [mock.call.start()] for idx, span in enumerate(trace): @@ -1266,11 +1266,13 @@ def _expected_llmobs_calls(trace, expected_spans_data: list): return expected_llmobs_writer_calls @staticmethod - def _expected_llmobs_chain_call(span, input_parameters=None): + def _expected_llmobs_chain_call(span, input_parameters=None, input_value=None, output_value=None): return _expected_llmobs_non_llm_span_event( span, span_kind="workflow", parameters=input_parameters, + input_value=input_value, + output_value=output_value, tags={ "ml_app": "langchain_community_test", }, @@ -1317,7 +1319,41 @@ def _expected_llmobs_llm_calls(span, provider="openai", input_role=None, output_ ) @classmethod - def _test_llmobs_invoke( + def _test_llmobs_llm_invoke( + cls, + provider, + generate_trace, + request_vcr, + mock_llmobs_writer, + mock_tracer, + cassette_name, + input_role=None, + output_role=None, + ): + LLMObs.disable() + LLMObs.enable(tracer=mock_tracer) + + with request_vcr.use_cassette(cassette_name): + generate_trace("Can you explain what an LLM chain is?") + span = mock_tracer.pop_traces()[0][0] + + expected_llmons_writer_calls = [ + mock.call.start(), + mock.call.enqueue( + cls._expected_llmobs_llm_calls( + span, + provider=provider, + input_role=input_role, + output_role=output_role, + ) + ), + ] + + assert mock_llmobs_writer.enqueue.call_count == 1 + mock_llmobs_writer.assert_has_calls(expected_llmons_writer_calls) + + @classmethod + def _test_llmobs_chain_invoke( cls, generate_trace, request_vcr, @@ -1334,7 +1370,9 @@ def _test_llmobs_invoke( generate_trace("Can you explain what an LLM chain is?") trace = mock_tracer.pop_traces()[0] - expected_llmobs_writer_calls = cls._expected_llmobs_calls(trace=trace, expected_spans_data=expected_spans_data) + expected_llmobs_writer_calls = cls._expected_llmobs_chain_calls( + trace=trace, expected_spans_data=expected_spans_data + ) assert mock_llmobs_writer.enqueue.call_count == len(expected_spans_data) mock_llmobs_writer.assert_has_calls(expected_llmobs_writer_calls) @@ -1342,62 +1380,67 @@ def _test_llmobs_invoke( def test_llmobs_openai_llm(self, langchain_openai, mock_llmobs_writer, mock_tracer, request_vcr): llm = langchain_openai.OpenAI() - self._test_llmobs_invoke( + self._test_llmobs_llm_invoke( generate_trace=llm.invoke, request_vcr=request_vcr, mock_llmobs_writer=mock_llmobs_writer, mock_tracer=mock_tracer, cassette_name="openai_completion_sync.yaml", + provider="openai", ) def test_llmobs_cohere_llm(self, langchain_community, mock_llmobs_writer, mock_tracer, request_vcr): llm = langchain_community.llms.Cohere(model="cohere.command-light-text-v14") - self._test_llmobs_invoke( + self._test_llmobs_llm_invoke( generate_trace=llm.invoke, request_vcr=request_vcr, mock_llmobs_writer=mock_llmobs_writer, mock_tracer=mock_tracer, cassette_name="cohere_completion_sync.yaml", - expected_spans_data=[("llm", {"provider": "cohere", "input_role": None, "output_role": None})], + provider="cohere", ) def test_llmobs_ai21_llm(self, langchain_community, mock_llmobs_writer, mock_tracer, request_vcr): llm = langchain_community.llms.AI21() - self._test_llmobs_invoke( + self._test_llmobs_llm_invoke( generate_trace=llm.invoke, request_vcr=request_vcr, mock_llmobs_writer=mock_llmobs_writer, mock_tracer=mock_tracer, cassette_name="ai21_completion_sync.yaml", - expected_spans_data=[("llm", {"provider": "ai21", "input_role": None, "output_role": None})], + provider="ai21", ) @flaky(1735812000) def test_llmobs_openai_chat_model(self, langchain_openai, mock_llmobs_writer, mock_tracer, request_vcr): chat = langchain_openai.ChatOpenAI(temperature=0, max_tokens=256) - self._test_llmobs_invoke( + self._test_llmobs_llm_invoke( generate_trace=lambda prompt: chat.invoke([langchain.schema.HumanMessage(content=prompt)]), request_vcr=request_vcr, mock_llmobs_writer=mock_llmobs_writer, mock_tracer=mock_tracer, cassette_name="openai_chat_completion_sync_call.yaml", - expected_spans_data=[("llm", {"provider": "openai", "input_role": "user", "output_role": "assistant"})], + provider="openai", + input_role="user", + output_role="assistant", ) @flaky(1735812000) def test_llmobs_openai_chat_model_custom_role(self, langchain_openai, mock_llmobs_writer, mock_tracer, request_vcr): chat = langchain_openai.ChatOpenAI(temperature=0, max_tokens=256) - self._test_llmobs_invoke( + self._test_llmobs_llm_invoke( generate_trace=lambda prompt: chat.invoke([langchain.schema.ChatMessage(content=prompt, role="custom")]), request_vcr=request_vcr, mock_llmobs_writer=mock_llmobs_writer, mock_tracer=mock_tracer, cassette_name="openai_chat_completion_sync_call.yaml", - expected_spans_data=[("llm", {"provider": "openai", "input_role": "custom", "output_role": "assistant"})], + provider="openai", + input_role="custom", + output_role="assistant", ) def test_llmobs_chain(self, langchain_core, langchain_openai, mock_llmobs_writer, mock_tracer, request_vcr): @@ -1408,14 +1451,20 @@ def test_llmobs_chain(self, langchain_core, langchain_openai, mock_llmobs_writer chain = prompt | llm - self._test_llmobs_invoke( + self._test_llmobs_chain_invoke( generate_trace=lambda prompt: chain.invoke({"input": prompt}), request_vcr=request_vcr, mock_llmobs_writer=mock_llmobs_writer, mock_tracer=mock_tracer, cassette_name="lcel_openai_chain_call.yaml", expected_spans_data=[ - ("chain", {"input_parameters": {"input": "Can you explain what an LLM chain is?"}}), + ( + "chain", + { + "input_value": str({"input": "Can you explain what an LLM chain is?"}), + "output_value": mock.ANY, + }, + ), ("llm", {"provider": "openai", "input_role": None, "output_role": None}), ], ) @@ -1433,7 +1482,7 @@ def test_llmobs_chain_nested(self, langchain_core, langchain_openai, mock_llmobs complete_chain = {"city": chain1, "language": itemgetter("language")} | chain2 - self._test_llmobs_invoke( + self._test_llmobs_chain_invoke( generate_trace=lambda inputs: complete_chain.invoke( {"person": "Spongebob Squarepants", "language": "Spanish"} ), @@ -1442,8 +1491,20 @@ def test_llmobs_chain_nested(self, langchain_core, langchain_openai, mock_llmobs mock_tracer=mock_tracer, cassette_name="lcel_openai_chain_nested.yaml", expected_spans_data=[ - ("chain", {"input_parameters": {"person": "Spongebob Squarepants", "language": "Spanish"}}), - ("chain", {"input_parameters": {"person": "Spongebob Squarepants", "language": "Spanish"}}), + ( + "chain", + { + "input_value": str({"person": "Spongebob Squarepants", "language": "Spanish"}), + "output_value": mock.ANY, + }, + ), + ( + "chain", + { + "input_value": str({"person": "Spongebob Squarepants", "language": "Spanish"}), + "output_value": mock.ANY, + }, + ), ("llm", {"provider": "openai", "input_role": "user", "output_role": "assistant"}), ("llm", {"provider": "openai", "input_role": "user", "output_role": "assistant"}), ], @@ -1456,14 +1517,20 @@ def test_llmobs_chain_batch(self, langchain_core, langchain_openai, mock_llmobs_ model = langchain_openai.ChatOpenAI() chain = {"topic": langchain_core.runnables.RunnablePassthrough()} | prompt | model | output_parser - self._test_llmobs_invoke( + self._test_llmobs_chain_invoke( generate_trace=lambda inputs: chain.batch(["chickens", "cows", "pigs"]), request_vcr=request_vcr, mock_llmobs_writer=mock_llmobs_writer, mock_tracer=mock_tracer, cassette_name="lcel_openai_chain_batch.yaml", expected_spans_data=[ - ("chain", {"input_parameters": {"input.0": "chickens", "input.1": "cows", "input.2": "pigs"}}), + ( + "chain", + { + "input_value": str({"input.0": "chickens", "input.1": "cows", "input.2": "pigs"}), + "output_value": mock.ANY, + }, + ), ("llm", {"provider": "openai", "input_role": "user", "output_role": "assistant"}), ("llm", {"provider": "openai", "input_role": "user", "output_role": "assistant"}), ("llm", {"provider": "openai", "input_role": "user", "output_role": "assistant"}), diff --git a/tests/llmobs/_utils.py b/tests/llmobs/_utils.py index 01b9f14c249..669b13ed87b 100644 --- a/tests/llmobs/_utils.py +++ b/tests/llmobs/_utils.py @@ -124,7 +124,7 @@ def _expected_llmobs_non_llm_span_event( if parameters is not None: meta_dict["input"].update({"parameters": parameters}) if output_value is not None: - meta_dict["output"].update({"messages": output_value}) + meta_dict["output"].update({"value": output_value}) if not meta_dict["input"]: meta_dict.pop("input") if not meta_dict["output"]: