mirror of https://github.com/hwchase17/langchain
Add `Runnable.astream_log()` (#10374)
--------- Co-authored-by: Bagatur <baskaryan@gmail.com>pull/10901/head
Nuno Campos
11 months ago
committed by
GitHub
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