diff --git a/tests/mlmodel_langchain/test_chain.py b/tests/mlmodel_langchain/test_chain.py index 2f52f8550..9a8d2cc74 100644 --- a/tests/mlmodel_langchain/test_chain.py +++ b/tests/mlmodel_langchain/test_chain.py @@ -1642,90 +1642,6 @@ def test_async_langchain_chain_outside_transaction( loop.run_until_complete(getattr(runnable, call_function)(input_)) -@pytest.mark.parametrize( - "create_function,call_function,call_function_args,call_function_kwargs,expected_events", - ( - pytest.param( - create_structured_output_runnable, - "ainvoke", - ({"input": "Sally is 13"},), - {"config": {"tags": ["bar"], "metadata": {"id": "123"}}}, - chat_completion_recorded_events_runnable_invoke, - id="runnable_chain.ainvoke-with-args-and-kwargs", - ), - pytest.param( - create_structured_output_chain, - "ainvoke", - ({"input": "Sally is 13"},), - {"config": {"tags": ["bar"], "metadata": {"id": "123"}}, "return_only_outputs": True}, - chat_completion_recorded_events_invoke, - id="chain.ainvoke-with-args-and-kwargs", - ), - ), -) -def test_multiple_async_langchain_chain( - set_trace_info, - json_schema, - prompt, - chat_openai_client, - create_function, - call_function, - call_function_args, - call_function_kwargs, - expected_events, - loop, -): - call1 = events_with_context_attrs(expected_events.copy()) - call1[0][1]["request_id"] = "b1883d9d-10d6-4b67-a911-f72849704e92" - call1[1][1]["request_id"] = "b1883d9d-10d6-4b67-a911-f72849704e92" - call1[2][1]["request_id"] = "b1883d9d-10d6-4b67-a911-f72849704e92" - call2 = events_with_context_attrs(expected_events.copy()) - call2[0][1]["request_id"] = "a58aa0c0-c854-4657-9e7b-4cce442f3b61" - call2[1][1]["request_id"] = "a58aa0c0-c854-4657-9e7b-4cce442f3b61" - call2[2][1]["request_id"] = "a58aa0c0-c854-4657-9e7b-4cce442f3b61" - - @reset_core_stats_engine() - @validate_custom_events(call1 + call2) - # 3 langchain events and 5 openai events. - @validate_custom_event_count(count=16) - @validate_transaction_metrics( - name="test_chain:test_multiple_async_langchain_chain.._test", - scoped_metrics=[(f"Llm/chain/LangChain/{call_function}", 2)], - rollup_metrics=[(f"Llm/chain/LangChain/{call_function}", 2)], - custom_metrics=[(f"Supportability/Python/ML/LangChain/{langchain.__version__}", 1)], - background_task=True, - ) - @background_task() - def _test(): - with patch("langchain_core.callbacks.manager.uuid", autospec=True) as mock_uuid: - mock_uuid.uuid4.side_effect = [ - uuid.UUID("b1883d9d-10d6-4b67-a911-f72849704e92"), # first call - uuid.UUID("a58aa0c0-c854-4657-9e7b-4cce442f3b61"), - uuid.UUID("a58aa0c0-c854-4657-9e7b-4cce442f3b61"), # second call - uuid.UUID("a58aa0c0-c854-4657-9e7b-4cce442f3b63"), - uuid.UUID("b1883d9d-10d6-4b67-a911-f72849704e93"), - uuid.UUID("a58aa0c0-c854-4657-9e7b-4cce442f3b64"), - uuid.UUID("a58aa0c0-c854-4657-9e7b-4cce442f3b65"), - uuid.UUID("a58aa0c0-c854-4657-9e7b-4cce442f3b66"), - ] - set_trace_info() - add_custom_attribute("llm.conversation_id", "my-awesome-id") - add_custom_attribute("llm.foo", "bar") - add_custom_attribute("non_llm_attr", "python-agent") - - runnable = create_function(json_schema, chat_openai_client, prompt) - with WithLlmCustomAttributes({"context": "attr"}): - call1 = asyncio.ensure_future( - getattr(runnable, call_function)(*call_function_args, **call_function_kwargs), loop=loop - ) - call2 = asyncio.ensure_future( - getattr(runnable, call_function)(*call_function_args, **call_function_kwargs), loop=loop - ) - loop.run_until_complete(asyncio.gather(call1, call2)) - - _test() - - @reset_core_stats_engine() @validate_custom_events(recorded_events_retrieval_chain_response) @validate_custom_event_count(count=17) diff --git a/tests/mlmodel_langchain/test_tool.py b/tests/mlmodel_langchain/test_tool.py index 18882b87d..9ce8d7c2a 100644 --- a/tests/mlmodel_langchain/test_tool.py +++ b/tests/mlmodel_langchain/test_tool.py @@ -384,9 +384,7 @@ def test_langchain_tool_disabled_ai_monitoring_events_async(set_trace_info, sing def test_langchain_multiple_async_calls(set_trace_info, single_arg_tool, multi_arg_tool, loop): call1 = single_arg_tool_recorded_events.copy() - call1[0][1]["run_id"] = "b1883d9d-10d6-4b67-a911-f72849704e92" call2 = multi_arg_tool_recorded_events.copy() - call2[0][1]["run_id"] = "a58aa0c0-c854-4657-9e7b-4cce442f3b61" expected_events = call1 + call2 @reset_core_stats_engine() @@ -401,27 +399,15 @@ def test_langchain_multiple_async_calls(set_trace_info, single_arg_tool, multi_a def _test(): set_trace_info() - with patch("langchain_core.callbacks.manager.uuid", autospec=True) as mock_uuid: - mock_uuid.uuid4.side_effect = [ - uuid.UUID("b1883d9d-10d6-4b67-a911-f72849704e92"), # first call - uuid.UUID("a58aa0c0-c854-4657-9e7b-4cce442f3b61"), - uuid.UUID("a58aa0c0-c854-4657-9e7b-4cce442f3b61"), # second call - uuid.UUID("a58aa0c0-c854-4657-9e7b-4cce442f3b63"), - uuid.UUID("b1883d9d-10d6-4b67-a911-f72849704e93"), - uuid.UUID("a58aa0c0-c854-4657-9e7b-4cce442f3b64"), - uuid.UUID("a58aa0c0-c854-4657-9e7b-4cce442f3b65"), - uuid.UUID("a58aa0c0-c854-4657-9e7b-4cce442f3b66"), - ] - - loop.run_until_complete( - asyncio.gather( - single_arg_tool.arun({"query": "Python Agent"}), - multi_arg_tool.arun( - {"first_num": 53, "second_num": 28}, - tags=["python", "test_tags"], - metadata={"test": "langchain", "test_run": True}, - ), - ) + loop.run_until_complete( + asyncio.gather( + single_arg_tool.arun({"query": "Python Agent"}), + multi_arg_tool.arun( + {"first_num": 53, "second_num": 28}, + tags=["python", "test_tags"], + metadata={"test": "langchain", "test_run": True}, + ), ) + ) _test()