mirror of https://github.com/hwchase17/langchain
Add `Runnable.astream_log()` (#10374)
--------- Co-authored-by: Bagatur <baskaryan@gmail.com>pull/10901/head
parent
a1ade48e8f
commit
fcb5aba9f0
@ -0,0 +1,289 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import math
|
||||||
|
import threading
|
||||||
|
from typing import (
|
||||||
|
Any,
|
||||||
|
AsyncIterator,
|
||||||
|
Dict,
|
||||||
|
List,
|
||||||
|
Optional,
|
||||||
|
Sequence,
|
||||||
|
TypedDict,
|
||||||
|
Union,
|
||||||
|
)
|
||||||
|
from uuid import UUID
|
||||||
|
|
||||||
|
import jsonpatch
|
||||||
|
from anyio import create_memory_object_stream
|
||||||
|
|
||||||
|
from langchain.callbacks.tracers.base import BaseTracer
|
||||||
|
from langchain.callbacks.tracers.schemas import Run
|
||||||
|
from langchain.schema.output import ChatGenerationChunk, GenerationChunk
|
||||||
|
|
||||||
|
|
||||||
|
class LogEntry(TypedDict):
|
||||||
|
id: str
|
||||||
|
"""ID of the sub-run."""
|
||||||
|
name: str
|
||||||
|
"""Name of the object being run."""
|
||||||
|
type: str
|
||||||
|
"""Type of the object being run, eg. prompt, chain, llm, etc."""
|
||||||
|
tags: List[str]
|
||||||
|
"""List of tags for the run."""
|
||||||
|
metadata: Dict[str, Any]
|
||||||
|
"""Key-value pairs of metadata for the run."""
|
||||||
|
start_time: str
|
||||||
|
"""ISO-8601 timestamp of when the run started."""
|
||||||
|
|
||||||
|
streamed_output_str: List[str]
|
||||||
|
"""List of LLM tokens streamed by this run, if applicable."""
|
||||||
|
final_output: Optional[Any]
|
||||||
|
"""Final output of this run.
|
||||||
|
Only available after the run has finished successfully."""
|
||||||
|
end_time: Optional[str]
|
||||||
|
"""ISO-8601 timestamp of when the run ended.
|
||||||
|
Only available after the run has finished."""
|
||||||
|
|
||||||
|
|
||||||
|
class RunState(TypedDict):
|
||||||
|
id: str
|
||||||
|
"""ID of the run."""
|
||||||
|
streamed_output: List[Any]
|
||||||
|
"""List of output chunks streamed by Runnable.stream()"""
|
||||||
|
final_output: Optional[Any]
|
||||||
|
"""Final output of the run, usually the result of aggregating streamed_output.
|
||||||
|
Only available after the run has finished successfully."""
|
||||||
|
|
||||||
|
logs: list[LogEntry]
|
||||||
|
"""List of sub-runs contained in this run, if any, in the order they were started.
|
||||||
|
If filters were supplied, this list will contain only the runs that matched the
|
||||||
|
filters."""
|
||||||
|
|
||||||
|
|
||||||
|
class RunLogPatch:
|
||||||
|
ops: List[Dict[str, Any]]
|
||||||
|
"""List of jsonpatch operations, which describe how to create the run state
|
||||||
|
from an empty dict. This is the minimal representation of the log, designed to
|
||||||
|
be serialized as JSON and sent over the wire to reconstruct the log on the other
|
||||||
|
side. Reconstruction of the state can be done with any jsonpatch-compliant library,
|
||||||
|
see https://jsonpatch.com for more information."""
|
||||||
|
|
||||||
|
def __init__(self, *ops: Dict[str, Any]) -> None:
|
||||||
|
self.ops = list(ops)
|
||||||
|
|
||||||
|
def __add__(self, other: Union[RunLogPatch, Any]) -> RunLogPatch:
|
||||||
|
if type(other) == RunLogPatch:
|
||||||
|
ops = self.ops + other.ops
|
||||||
|
state = jsonpatch.apply_patch(None, ops)
|
||||||
|
return RunLog(*ops, state=state)
|
||||||
|
|
||||||
|
raise TypeError(
|
||||||
|
f"unsupported operand type(s) for +: '{type(self)}' and '{type(other)}'"
|
||||||
|
)
|
||||||
|
|
||||||
|
def __repr__(self) -> str:
|
||||||
|
from pprint import pformat
|
||||||
|
|
||||||
|
return f"RunLogPatch(ops={pformat(self.ops)})"
|
||||||
|
|
||||||
|
def __eq__(self, other: object) -> bool:
|
||||||
|
return isinstance(other, RunLogPatch) and self.ops == other.ops
|
||||||
|
|
||||||
|
|
||||||
|
class RunLog(RunLogPatch):
|
||||||
|
state: RunState
|
||||||
|
"""Current state of the log, obtained from applying all ops in sequence."""
|
||||||
|
|
||||||
|
def __init__(self, *ops: Dict[str, Any], state: RunState) -> None:
|
||||||
|
super().__init__(*ops)
|
||||||
|
self.state = state
|
||||||
|
|
||||||
|
def __add__(self, other: Union[RunLogPatch, Any]) -> RunLogPatch:
|
||||||
|
if type(other) == RunLogPatch:
|
||||||
|
ops = self.ops + other.ops
|
||||||
|
state = jsonpatch.apply_patch(self.state, other.ops)
|
||||||
|
return RunLog(*ops, state=state)
|
||||||
|
|
||||||
|
raise TypeError(
|
||||||
|
f"unsupported operand type(s) for +: '{type(self)}' and '{type(other)}'"
|
||||||
|
)
|
||||||
|
|
||||||
|
def __repr__(self) -> str:
|
||||||
|
from pprint import pformat
|
||||||
|
|
||||||
|
return f"RunLog(state={pformat(self.state)})"
|
||||||
|
|
||||||
|
|
||||||
|
class LogStreamCallbackHandler(BaseTracer):
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
*,
|
||||||
|
auto_close: bool = True,
|
||||||
|
include_names: Optional[Sequence[str]] = None,
|
||||||
|
include_types: Optional[Sequence[str]] = None,
|
||||||
|
include_tags: Optional[Sequence[str]] = None,
|
||||||
|
exclude_names: Optional[Sequence[str]] = None,
|
||||||
|
exclude_types: Optional[Sequence[str]] = None,
|
||||||
|
exclude_tags: Optional[Sequence[str]] = None,
|
||||||
|
) -> None:
|
||||||
|
super().__init__()
|
||||||
|
|
||||||
|
self.auto_close = auto_close
|
||||||
|
self.include_names = include_names
|
||||||
|
self.include_types = include_types
|
||||||
|
self.include_tags = include_tags
|
||||||
|
self.exclude_names = exclude_names
|
||||||
|
self.exclude_types = exclude_types
|
||||||
|
self.exclude_tags = exclude_tags
|
||||||
|
|
||||||
|
send_stream, receive_stream = create_memory_object_stream(
|
||||||
|
math.inf, item_type=RunLogPatch
|
||||||
|
)
|
||||||
|
self.lock = threading.Lock()
|
||||||
|
self.send_stream = send_stream
|
||||||
|
self.receive_stream = receive_stream
|
||||||
|
self._index_map: Dict[UUID, int] = {}
|
||||||
|
|
||||||
|
def __aiter__(self) -> AsyncIterator[RunLogPatch]:
|
||||||
|
return self.receive_stream.__aiter__()
|
||||||
|
|
||||||
|
def include_run(self, run: Run) -> bool:
|
||||||
|
if run.parent_run_id is None:
|
||||||
|
return False
|
||||||
|
|
||||||
|
run_tags = run.tags or []
|
||||||
|
|
||||||
|
if (
|
||||||
|
self.include_names is None
|
||||||
|
and self.include_types is None
|
||||||
|
and self.include_tags is None
|
||||||
|
):
|
||||||
|
include = True
|
||||||
|
else:
|
||||||
|
include = False
|
||||||
|
|
||||||
|
if self.include_names is not None:
|
||||||
|
include = include or run.name in self.include_names
|
||||||
|
if self.include_types is not None:
|
||||||
|
include = include or run.run_type in self.include_types
|
||||||
|
if self.include_tags is not None:
|
||||||
|
include = include or any(tag in self.include_tags for tag in run_tags)
|
||||||
|
|
||||||
|
if self.exclude_names is not None:
|
||||||
|
include = include and run.name not in self.exclude_names
|
||||||
|
if self.exclude_types is not None:
|
||||||
|
include = include and run.run_type not in self.exclude_types
|
||||||
|
if self.exclude_tags is not None:
|
||||||
|
include = include and all(tag not in self.exclude_tags for tag in run_tags)
|
||||||
|
|
||||||
|
return include
|
||||||
|
|
||||||
|
def _persist_run(self, run: Run) -> None:
|
||||||
|
# This is a legacy method only called once for an entire run tree
|
||||||
|
# therefore not useful here
|
||||||
|
pass
|
||||||
|
|
||||||
|
def _on_run_create(self, run: Run) -> None:
|
||||||
|
"""Start a run."""
|
||||||
|
if run.parent_run_id is None:
|
||||||
|
self.send_stream.send_nowait(
|
||||||
|
RunLogPatch(
|
||||||
|
{
|
||||||
|
"op": "replace",
|
||||||
|
"path": "",
|
||||||
|
"value": RunState(
|
||||||
|
id=run.id,
|
||||||
|
streamed_output=[],
|
||||||
|
final_output=None,
|
||||||
|
logs=[],
|
||||||
|
),
|
||||||
|
}
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
if not self.include_run(run):
|
||||||
|
return
|
||||||
|
|
||||||
|
# Determine previous index, increment by 1
|
||||||
|
with self.lock:
|
||||||
|
self._index_map[run.id] = max(self._index_map.values(), default=-1) + 1
|
||||||
|
|
||||||
|
# Add the run to the stream
|
||||||
|
self.send_stream.send_nowait(
|
||||||
|
RunLogPatch(
|
||||||
|
{
|
||||||
|
"op": "add",
|
||||||
|
"path": f"/logs/{self._index_map[run.id]}",
|
||||||
|
"value": LogEntry(
|
||||||
|
id=str(run.id),
|
||||||
|
name=run.name,
|
||||||
|
type=run.run_type,
|
||||||
|
tags=run.tags or [],
|
||||||
|
metadata=run.extra.get("metadata", {}),
|
||||||
|
start_time=run.start_time.isoformat(timespec="milliseconds"),
|
||||||
|
streamed_output_str=[],
|
||||||
|
final_output=None,
|
||||||
|
end_time=None,
|
||||||
|
),
|
||||||
|
}
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
def _on_run_update(self, run: Run) -> None:
|
||||||
|
"""Finish a run."""
|
||||||
|
try:
|
||||||
|
index = self._index_map.get(run.id)
|
||||||
|
|
||||||
|
if index is None:
|
||||||
|
return
|
||||||
|
|
||||||
|
self.send_stream.send_nowait(
|
||||||
|
RunLogPatch(
|
||||||
|
{
|
||||||
|
"op": "add",
|
||||||
|
"path": f"/logs/{index}/final_output",
|
||||||
|
"value": run.outputs,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"op": "add",
|
||||||
|
"path": f"/logs/{index}/end_time",
|
||||||
|
"value": run.end_time.isoformat(timespec="milliseconds"),
|
||||||
|
},
|
||||||
|
)
|
||||||
|
)
|
||||||
|
finally:
|
||||||
|
if run.parent_run_id is None:
|
||||||
|
self.send_stream.send_nowait(
|
||||||
|
RunLogPatch(
|
||||||
|
{
|
||||||
|
"op": "replace",
|
||||||
|
"path": "/final_output",
|
||||||
|
"value": run.outputs,
|
||||||
|
}
|
||||||
|
)
|
||||||
|
)
|
||||||
|
if self.auto_close:
|
||||||
|
self.send_stream.close()
|
||||||
|
|
||||||
|
def _on_llm_new_token(
|
||||||
|
self,
|
||||||
|
run: Run,
|
||||||
|
token: str,
|
||||||
|
chunk: Optional[Union[GenerationChunk, ChatGenerationChunk]],
|
||||||
|
) -> None:
|
||||||
|
"""Process new LLM token."""
|
||||||
|
index = self._index_map.get(run.id)
|
||||||
|
|
||||||
|
if index is None:
|
||||||
|
return
|
||||||
|
|
||||||
|
self.send_stream.send_nowait(
|
||||||
|
RunLogPatch(
|
||||||
|
{
|
||||||
|
"op": "add",
|
||||||
|
"path": f"/logs/{index}/streamed_output_str/-",
|
||||||
|
"value": token,
|
||||||
|
}
|
||||||
|
)
|
||||||
|
)
|
Loading…
Reference in New Issue