diff --git a/langfuse/langchain/CallbackHandler.py b/langfuse/langchain/CallbackHandler.py index 898a6b22c..65ce1b2d8 100644 --- a/langfuse/langchain/CallbackHandler.py +++ b/langfuse/langchain/CallbackHandler.py @@ -1369,6 +1369,8 @@ def on_llm_error( except Exception as e: langfuse_logger.exception(e) finally: + self._updated_completion_start_time_memo.discard(run_id) + if parent_run_id is None: self._reset(run_id) diff --git a/tests/unit/test_langchain.py b/tests/unit/test_langchain.py index 12639f0bb..dbbb3d93b 100644 --- a/tests/unit/test_langchain.py +++ b/tests/unit/test_langchain.py @@ -356,6 +356,51 @@ class DummyControlFlowError(RuntimeError): ) +def test_streamed_llm_error_cleans_only_its_completion_start_time_memo( + langfuse_memory_client, +): + handler = CallbackHandler() + failed_run_id = uuid4() + active_run_id = uuid4() + failed_context = copy_context() + active_context = copy_context() + + for context, run_id in ( + (failed_context, failed_run_id), + (active_context, active_run_id), + ): + context.run( + handler.on_llm_start, + {"name": "TestLLM"}, + ["prompt"], + run_id=run_id, + metadata={"ls_model_name": "test-model"}, + invocation_params={}, + ) + context.run(handler.on_llm_new_token, "first token", run_id=run_id) + + response = LLMResult(generations=[[Generation(text="done")]], llm_output={}) + + try: + assert failed_run_id in handler._updated_completion_start_time_memo + assert active_run_id in handler._updated_completion_start_time_memo + + failed_context.run( + handler.on_llm_error, + RuntimeError("synthetic failure"), + run_id=failed_run_id, + ) + + assert active_run_id in handler._updated_completion_start_time_memo + assert failed_run_id not in handler._updated_completion_start_time_memo + finally: + active_context.run( + handler.on_llm_end, + response, + run_id=active_run_id, + ) + + def test_control_flow_resume_uses_thread_keyed_explicit_resume_context( memory_exporter, langfuse_memory_client, monkeypatch ):