diff --git a/langchain/callbacks/tracers/base.py b/langchain/callbacks/tracers/base.py index edee1e2799..3322fbffc7 100644 --- a/langchain/callbacks/tracers/base.py +++ b/langchain/callbacks/tracers/base.py @@ -92,13 +92,15 @@ class BaseTracer(BaseCallbackHandler, ABC): """Start a trace for an LLM run.""" parent_run_id_ = str(parent_run_id) if parent_run_id else None execution_order = self._get_execution_order(parent_run_id_) + start_time = datetime.utcnow() llm_run = Run( id=run_id, parent_run_id=parent_run_id, serialized=serialized, inputs={"prompts": prompts}, extra=kwargs, - start_time=datetime.utcnow(), + events=[{"name": "start", "time": start_time}], + start_time=start_time, execution_order=execution_order, child_execution_order=execution_order, run_type=RunTypeEnum.llm, @@ -107,6 +109,30 @@ class BaseTracer(BaseCallbackHandler, ABC): self._start_trace(llm_run) self._on_llm_start(llm_run) + def on_llm_new_token( + self, + token: str, + *, + run_id: UUID, + parent_run_id: Optional[UUID] = None, + **kwargs: Any, + ) -> None: + """Run on new LLM token. Only available when streaming is enabled.""" + if not run_id: + raise TracerException("No run_id provided for on_llm_new_token callback.") + + run_id_ = str(run_id) + llm_run = self.run_map.get(run_id_) + if llm_run is None or llm_run.run_type != RunTypeEnum.llm: + raise TracerException("No LLM Run found to be traced") + llm_run.events.append( + { + "name": "new_token", + "time": datetime.utcnow(), + "kwargs": {"token": token}, + }, + ) + def on_llm_end(self, response: LLMResult, *, run_id: UUID, **kwargs: Any) -> None: """End a trace for an LLM run.""" if not run_id: @@ -118,6 +144,7 @@ class BaseTracer(BaseCallbackHandler, ABC): raise TracerException("No LLM Run found to be traced") llm_run.outputs = response.dict() llm_run.end_time = datetime.utcnow() + llm_run.events.append({"name": "end", "time": llm_run.end_time}) self._end_trace(llm_run) self._on_llm_end(llm_run) @@ -138,6 +165,7 @@ class BaseTracer(BaseCallbackHandler, ABC): raise TracerException("No LLM Run found to be traced") llm_run.error = repr(error) llm_run.end_time = datetime.utcnow() + llm_run.events.append({"name": "error", "time": llm_run.end_time}) self._end_trace(llm_run) self._on_chain_error(llm_run) @@ -154,13 +182,15 @@ class BaseTracer(BaseCallbackHandler, ABC): """Start a trace for a chain run.""" parent_run_id_ = str(parent_run_id) if parent_run_id else None execution_order = self._get_execution_order(parent_run_id_) + start_time = datetime.utcnow() chain_run = Run( id=run_id, parent_run_id=parent_run_id, serialized=serialized, inputs=inputs, extra=kwargs, - start_time=datetime.utcnow(), + events=[{"name": "start", "time": start_time}], + start_time=start_time, execution_order=execution_order, child_execution_order=execution_order, child_runs=[], @@ -182,6 +212,7 @@ class BaseTracer(BaseCallbackHandler, ABC): chain_run.outputs = outputs chain_run.end_time = datetime.utcnow() + chain_run.events.append({"name": "end", "time": chain_run.end_time}) self._end_trace(chain_run) self._on_chain_end(chain_run) @@ -201,6 +232,7 @@ class BaseTracer(BaseCallbackHandler, ABC): chain_run.error = repr(error) chain_run.end_time = datetime.utcnow() + chain_run.events.append({"name": "error", "time": chain_run.end_time}) self._end_trace(chain_run) self._on_chain_error(chain_run) @@ -217,13 +249,15 @@ class BaseTracer(BaseCallbackHandler, ABC): """Start a trace for a tool run.""" parent_run_id_ = str(parent_run_id) if parent_run_id else None execution_order = self._get_execution_order(parent_run_id_) + start_time = datetime.utcnow() tool_run = Run( id=run_id, parent_run_id=parent_run_id, serialized=serialized, inputs={"input": input_str}, extra=kwargs, - start_time=datetime.utcnow(), + events=[{"name": "start", "time": start_time}], + start_time=start_time, execution_order=execution_order, child_execution_order=execution_order, child_runs=[], @@ -243,6 +277,7 @@ class BaseTracer(BaseCallbackHandler, ABC): tool_run.outputs = {"output": output} tool_run.end_time = datetime.utcnow() + tool_run.events.append({"name": "end", "time": tool_run.end_time}) self._end_trace(tool_run) self._on_tool_end(tool_run) @@ -262,6 +297,7 @@ class BaseTracer(BaseCallbackHandler, ABC): tool_run.error = repr(error) tool_run.end_time = datetime.utcnow() + tool_run.events.append({"name": "error", "time": tool_run.end_time}) self._end_trace(tool_run) self._on_tool_error(tool_run) @@ -276,13 +312,15 @@ class BaseTracer(BaseCallbackHandler, ABC): """Run when Retriever starts running.""" parent_run_id_ = str(parent_run_id) if parent_run_id else None execution_order = self._get_execution_order(parent_run_id_) + start_time = datetime.utcnow() retrieval_run = Run( id=run_id, name="Retriever", parent_run_id=parent_run_id, inputs={"query": query}, extra=kwargs, - start_time=datetime.utcnow(), + events=[{"name": "start", "time": start_time}], + start_time=start_time, execution_order=execution_order, child_execution_order=execution_order, child_runs=[], @@ -307,6 +345,7 @@ class BaseTracer(BaseCallbackHandler, ABC): retrieval_run.error = repr(error) retrieval_run.end_time = datetime.utcnow() + retrieval_run.events.append({"name": "error", "time": retrieval_run.end_time}) self._end_trace(retrieval_run) self._on_retriever_error(retrieval_run) @@ -321,6 +360,7 @@ class BaseTracer(BaseCallbackHandler, ABC): raise TracerException("No retriever Run found to be traced") retrieval_run.outputs = {"documents": documents} retrieval_run.end_time = datetime.utcnow() + retrieval_run.events.append({"name": "end", "time": retrieval_run.end_time}) self._end_trace(retrieval_run) self._on_retriever_end(retrieval_run) diff --git a/langchain/callbacks/tracers/langchain.py b/langchain/callbacks/tracers/langchain.py index 1c0095dd4d..f42f736381 100644 --- a/langchain/callbacks/tracers/langchain.py +++ b/langchain/callbacks/tracers/langchain.py @@ -78,13 +78,15 @@ class LangChainTracer(BaseTracer): """Start a trace for an LLM run.""" parent_run_id_ = str(parent_run_id) if parent_run_id else None execution_order = self._get_execution_order(parent_run_id_) + start_time = datetime.utcnow() chat_model_run = Run( id=run_id, parent_run_id=parent_run_id, serialized=serialized, inputs={"messages": [messages_to_dict(batch) for batch in messages]}, extra=kwargs, - start_time=datetime.utcnow(), + events=[{"name": "start", "time": start_time}], + start_time=start_time, execution_order=execution_order, child_execution_order=execution_order, run_type=RunTypeEnum.llm, diff --git a/tests/unit_tests/callbacks/tracers/test_base_tracer.py b/tests/unit_tests/callbacks/tracers/test_base_tracer.py index c73253a2b9..c4b03b45bd 100644 --- a/tests/unit_tests/callbacks/tracers/test_base_tracer.py +++ b/tests/unit_tests/callbacks/tracers/test_base_tracer.py @@ -40,6 +40,10 @@ def test_tracer_llm_run() -> None: parent_run_id=None, start_time=datetime.utcnow(), end_time=datetime.utcnow(), + events=[ + {"name": "start", "time": datetime.utcnow()}, + {"name": "end", "time": datetime.utcnow()}, + ], extra={}, execution_order=1, child_execution_order=1, @@ -69,6 +73,10 @@ def test_tracer_chat_model_run() -> None: name="chat_model", start_time=datetime.utcnow(), end_time=datetime.utcnow(), + events=[ + {"name": "start", "time": datetime.utcnow()}, + {"name": "end", "time": datetime.utcnow()}, + ], extra={}, execution_order=1, child_execution_order=1, @@ -101,6 +109,10 @@ def test_tracer_multiple_llm_runs() -> None: name="llm", start_time=datetime.utcnow(), end_time=datetime.utcnow(), + events=[ + {"name": "start", "time": datetime.utcnow()}, + {"name": "end", "time": datetime.utcnow()}, + ], extra={}, execution_order=1, child_execution_order=1, @@ -128,6 +140,10 @@ def test_tracer_chain_run() -> None: id=str(uuid), start_time=datetime.utcnow(), end_time=datetime.utcnow(), + events=[ + {"name": "start", "time": datetime.utcnow()}, + {"name": "end", "time": datetime.utcnow()}, + ], extra={}, execution_order=1, child_execution_order=1, @@ -152,6 +168,10 @@ def test_tracer_tool_run() -> None: id=str(uuid), start_time=datetime.utcnow(), end_time=datetime.utcnow(), + events=[ + {"name": "start", "time": datetime.utcnow()}, + {"name": "end", "time": datetime.utcnow()}, + ], extra={}, execution_order=1, child_execution_order=1, @@ -208,6 +228,10 @@ def test_tracer_nested_run() -> None: error=None, start_time=datetime.utcnow(), end_time=datetime.utcnow(), + events=[ + {"name": "start", "time": datetime.utcnow()}, + {"name": "end", "time": datetime.utcnow()}, + ], extra={}, execution_order=1, child_execution_order=4, @@ -221,6 +245,10 @@ def test_tracer_nested_run() -> None: parent_run_id=chain_uuid, start_time=datetime.utcnow(), end_time=datetime.utcnow(), + events=[ + {"name": "start", "time": datetime.utcnow()}, + {"name": "end", "time": datetime.utcnow()}, + ], extra={}, execution_order=2, child_execution_order=3, @@ -236,6 +264,10 @@ def test_tracer_nested_run() -> None: error=None, start_time=datetime.utcnow(), end_time=datetime.utcnow(), + events=[ + {"name": "start", "time": datetime.utcnow()}, + {"name": "end", "time": datetime.utcnow()}, + ], extra={}, execution_order=3, child_execution_order=3, @@ -252,6 +284,10 @@ def test_tracer_nested_run() -> None: error=None, start_time=datetime.utcnow(), end_time=datetime.utcnow(), + events=[ + {"name": "start", "time": datetime.utcnow()}, + {"name": "end", "time": datetime.utcnow()}, + ], extra={}, execution_order=4, child_execution_order=4, @@ -276,6 +312,10 @@ def test_tracer_llm_run_on_error() -> None: id=str(uuid), start_time=datetime.utcnow(), end_time=datetime.utcnow(), + events=[ + {"name": "start", "time": datetime.utcnow()}, + {"name": "error", "time": datetime.utcnow()}, + ], extra={}, execution_order=1, child_execution_order=1, @@ -302,6 +342,10 @@ def test_tracer_chain_run_on_error() -> None: id=str(uuid), start_time=datetime.utcnow(), end_time=datetime.utcnow(), + events=[ + {"name": "start", "time": datetime.utcnow()}, + {"name": "error", "time": datetime.utcnow()}, + ], extra={}, execution_order=1, child_execution_order=1, @@ -328,6 +372,10 @@ def test_tracer_tool_run_on_error() -> None: id=str(uuid), start_time=datetime.utcnow(), end_time=datetime.utcnow(), + events=[ + {"name": "start", "time": datetime.utcnow()}, + {"name": "error", "time": datetime.utcnow()}, + ], extra={}, execution_order=1, child_execution_order=1, @@ -395,6 +443,10 @@ def test_tracer_nested_runs_on_error() -> None: id=str(chain_uuid), start_time=datetime.utcnow(), end_time=datetime.utcnow(), + events=[ + {"name": "start", "time": datetime.utcnow()}, + {"name": "error", "time": datetime.utcnow()}, + ], extra={}, execution_order=1, child_execution_order=5, @@ -409,6 +461,10 @@ def test_tracer_nested_runs_on_error() -> None: parent_run_id=str(chain_uuid), start_time=datetime.utcnow(), end_time=datetime.utcnow(), + events=[ + {"name": "start", "time": datetime.utcnow()}, + {"name": "end", "time": datetime.utcnow()}, + ], extra={}, execution_order=2, child_execution_order=2, @@ -423,6 +479,10 @@ def test_tracer_nested_runs_on_error() -> None: parent_run_id=str(chain_uuid), start_time=datetime.utcnow(), end_time=datetime.utcnow(), + events=[ + {"name": "start", "time": datetime.utcnow()}, + {"name": "end", "time": datetime.utcnow()}, + ], extra={}, execution_order=3, child_execution_order=3, @@ -437,6 +497,10 @@ def test_tracer_nested_runs_on_error() -> None: parent_run_id=str(chain_uuid), start_time=datetime.utcnow(), end_time=datetime.utcnow(), + events=[ + {"name": "start", "time": datetime.utcnow()}, + {"name": "error", "time": datetime.utcnow()}, + ], extra={}, execution_order=4, child_execution_order=5, @@ -451,6 +515,10 @@ def test_tracer_nested_runs_on_error() -> None: parent_run_id=str(tool_uuid), start_time=datetime.utcnow(), end_time=datetime.utcnow(), + events=[ + {"name": "start", "time": datetime.utcnow()}, + {"name": "error", "time": datetime.utcnow()}, + ], extra={}, execution_order=5, child_execution_order=5,