Add events to tracer runs (#7090)

<!-- Thank you for contributing to LangChain!

Replace this comment with:
  - Description: a description of the change, 
  - Issue: the issue # it fixes (if applicable),
  - Dependencies: any dependencies required for this change,
- Tag maintainer: for a quicker response, tag the relevant maintainer
(see below),
- Twitter handle: we announce bigger features on Twitter. If your PR
gets announced and you'd like a mention, we'll gladly shout you out!

If you're adding a new integration, please include:
1. a test for the integration, preferably unit tests that do not rely on
network access,
  2. an example notebook showing its use.

Maintainer responsibilities:
  - General / Misc / if you don't know who to tag: @dev2049
  - DataLoaders / VectorStores / Retrievers: @rlancemartin, @eyurtsev
  - Models / Prompts: @hwchase17, @dev2049
  - Memory: @hwchase17
  - Agents / Tools / Toolkits: @vowelparrot
  - Tracing / Callbacks: @agola11
  - Async: @agola11

If no one reviews your PR within a few days, feel free to @-mention the
same people again.

See contribution guidelines for more information on how to write/run
tests, lint, etc:
https://github.com/hwchase17/langchain/blob/master/.github/CONTRIBUTING.md
 -->
This commit is contained in:
Nuno Campos 2023-07-03 20:43:43 +01:00 committed by GitHub
parent e49abd1277
commit c8f8b1b327
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 115 additions and 5 deletions

View File

@ -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)

View File

@ -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,

View File

@ -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,