aerograph-langchain 0.2.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- aerograph_langchain-0.2.0/.gitignore +37 -0
- aerograph_langchain-0.2.0/PKG-INFO +117 -0
- aerograph_langchain-0.2.0/README.md +88 -0
- aerograph_langchain-0.2.0/examples/langchain_demo.py +53 -0
- aerograph_langchain-0.2.0/pyproject.toml +56 -0
- aerograph_langchain-0.2.0/src/aerograph_langchain/__init__.py +5 -0
- aerograph_langchain-0.2.0/src/aerograph_langchain/handler.py +252 -0
- aerograph_langchain-0.2.0/src/aerograph_langchain/langgraph.py +56 -0
- aerograph_langchain-0.2.0/src/aerograph_langchain/mapping.py +174 -0
- aerograph_langchain-0.2.0/src/aerograph_langchain/retriever.py +56 -0
- aerograph_langchain-0.2.0/src/aerograph_langchain/span_ids.py +26 -0
- aerograph_langchain-0.2.0/src/aerograph_langchain/streaming.py +51 -0
- aerograph_langchain-0.2.0/tests/fixtures/langchain_run.json +90 -0
- aerograph_langchain-0.2.0/tests/integration/test_end_to_end.py +43 -0
- aerograph_langchain-0.2.0/tests/test_error_handling.py +56 -0
- aerograph_langchain-0.2.0/tests/test_handler_mapping.py +125 -0
- aerograph_langchain-0.2.0/tests/test_langgraph.py +57 -0
- aerograph_langchain-0.2.0/tests/test_retriever.py +37 -0
- aerograph_langchain-0.2.0/tests/test_span_id_derivation.py +38 -0
- aerograph_langchain-0.2.0/tests/test_streaming.py +40 -0
|
@@ -0,0 +1,37 @@
|
|
|
1
|
+
node_modules/
|
|
2
|
+
dist/
|
|
3
|
+
build/
|
|
4
|
+
.vite/
|
|
5
|
+
coverage/
|
|
6
|
+
*.log
|
|
7
|
+
.env*
|
|
8
|
+
.specify/
|
|
9
|
+
apps/collector/data/
|
|
10
|
+
.DS_Store
|
|
11
|
+
.github/agents
|
|
12
|
+
.github/prompts
|
|
13
|
+
.github/copilot-instructions.md
|
|
14
|
+
Thumbs.db
|
|
15
|
+
.agents/
|
|
16
|
+
apps/collector/data/
|
|
17
|
+
|
|
18
|
+
# Python
|
|
19
|
+
__pycache__/
|
|
20
|
+
*.py[cod]
|
|
21
|
+
.venv/
|
|
22
|
+
venv/
|
|
23
|
+
*.egg-info/
|
|
24
|
+
.pytest_cache/
|
|
25
|
+
.mypy_cache/
|
|
26
|
+
.ruff_cache/
|
|
27
|
+
.DS_Store
|
|
28
|
+
Thumbs.db
|
|
29
|
+
*.tmp
|
|
30
|
+
*.swp
|
|
31
|
+
.vscode/
|
|
32
|
+
.idea/
|
|
33
|
+
specs/
|
|
34
|
+
docs/architecture
|
|
35
|
+
.vscode/
|
|
36
|
+
AGENTS.md
|
|
37
|
+
uv.lock
|
|
@@ -0,0 +1,117 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: aerograph-langchain
|
|
3
|
+
Version: 0.2.0
|
|
4
|
+
Summary: LangChain callback adapter for AeroGraph — automatically record LangChain traces.
|
|
5
|
+
Project-URL: Homepage, https://github.com/SGcpu/AeroGraph
|
|
6
|
+
Project-URL: Repository, https://github.com/SGcpu/AeroGraph
|
|
7
|
+
Project-URL: Issues, https://github.com/SGcpu/AeroGraph/issues
|
|
8
|
+
License: Apache-2.0
|
|
9
|
+
Keywords: aerograph,agent,ai,langchain,tracing
|
|
10
|
+
Classifier: Development Status :: 3 - Alpha
|
|
11
|
+
Classifier: Intended Audience :: Developers
|
|
12
|
+
Classifier: License :: OSI Approved :: Apache Software License
|
|
13
|
+
Classifier: Programming Language :: Python :: 3
|
|
14
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
15
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
16
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
17
|
+
Classifier: Topic :: Software Development :: Libraries
|
|
18
|
+
Requires-Python: >=3.10
|
|
19
|
+
Requires-Dist: aerograph-sdk>=0.2.0
|
|
20
|
+
Requires-Dist: langchain-core>=0.2.0
|
|
21
|
+
Provides-Extra: dev
|
|
22
|
+
Requires-Dist: langchain-core>=0.2.0; extra == 'dev'
|
|
23
|
+
Requires-Dist: langchain>=0.2.0; extra == 'dev'
|
|
24
|
+
Requires-Dist: pytest-asyncio>=0.24.0; extra == 'dev'
|
|
25
|
+
Requires-Dist: pytest>=8.0.0; extra == 'dev'
|
|
26
|
+
Provides-Extra: langgraph
|
|
27
|
+
Requires-Dist: langgraph>=0.1.0; extra == 'langgraph'
|
|
28
|
+
Description-Content-Type: text/markdown
|
|
29
|
+
|
|
30
|
+
# aerograph-langchain
|
|
31
|
+
|
|
32
|
+
LangChain callback adapter for [AeroGraph](https://github.com/SGcpu/AeroGraph).
|
|
33
|
+
|
|
34
|
+
Automatically records LangChain chain, LLM, tool, and retriever events as AeroGraph trace events.
|
|
35
|
+
|
|
36
|
+
## Installation
|
|
37
|
+
|
|
38
|
+
When published, you can install the adapter via pip:
|
|
39
|
+
|
|
40
|
+
```bash
|
|
41
|
+
pip install aerograph-langchain
|
|
42
|
+
```
|
|
43
|
+
|
|
44
|
+
For local testing or development, install the packages in editable mode from the repository root:
|
|
45
|
+
|
|
46
|
+
```bash
|
|
47
|
+
# Install the core SDK first
|
|
48
|
+
pip install -e python/aerograph-sdk
|
|
49
|
+
|
|
50
|
+
# Install the LangChain adapter
|
|
51
|
+
pip install -e python/aerograph-langchain
|
|
52
|
+
```
|
|
53
|
+
|
|
54
|
+
## Usage
|
|
55
|
+
|
|
56
|
+
Integrating `AeroGraphCallbackHandler` into your existing LangChain codebase is simple. Setup a `FlightRecorder`, initialize the handler, and pass it to your chain or model invocations.
|
|
57
|
+
|
|
58
|
+
### Basic Example
|
|
59
|
+
|
|
60
|
+
```python
|
|
61
|
+
import asyncio
|
|
62
|
+
from langchain_core.messages import HumanMessage
|
|
63
|
+
from langchain_openai import ChatOpenAI
|
|
64
|
+
from aerograph_sdk.recorder import FlightRecorder
|
|
65
|
+
from aerograph_langchain.handler import AeroGraphCallbackHandler
|
|
66
|
+
|
|
67
|
+
async def main():
|
|
68
|
+
# 1. Initialize the FlightRecorder pointing to your collector
|
|
69
|
+
recorder = FlightRecorder(
|
|
70
|
+
endpoint="http://localhost:4317",
|
|
71
|
+
actor={"id": "my-agent", "name": "TravelPlanner"}
|
|
72
|
+
)
|
|
73
|
+
|
|
74
|
+
# 2. Initialize the AeroGraph callback handler
|
|
75
|
+
handler = AeroGraphCallbackHandler(recorder)
|
|
76
|
+
|
|
77
|
+
# 3. Create your LangChain chat model
|
|
78
|
+
model = ChatOpenAI(model="gpt-4o")
|
|
79
|
+
|
|
80
|
+
# 4. Invoke the model and pass the handler in callbacks
|
|
81
|
+
response = await model.ainvoke(
|
|
82
|
+
[HumanMessage(content="What are 3 fun things to do in San Francisco?")],
|
|
83
|
+
config={"callbacks": [handler]}
|
|
84
|
+
)
|
|
85
|
+
print(response.content)
|
|
86
|
+
|
|
87
|
+
if __name__ == "__main__":
|
|
88
|
+
asyncio.run(main())
|
|
89
|
+
```
|
|
90
|
+
|
|
91
|
+
### Advanced Usage with Chains and RAG
|
|
92
|
+
|
|
93
|
+
You can attach the callback handler at the chain execution level. LangChain automatically propagates the callbacks down to all sub-chains, LLMs, retrievers, and tool invocations.
|
|
94
|
+
|
|
95
|
+
```python
|
|
96
|
+
# Pass the handler to the chain invoke call
|
|
97
|
+
result = rag_chain.invoke(
|
|
98
|
+
"How do I configure the server?",
|
|
99
|
+
config={"callbacks": [handler]}
|
|
100
|
+
)
|
|
101
|
+
```
|
|
102
|
+
|
|
103
|
+
## Features and Event Mapping
|
|
104
|
+
|
|
105
|
+
The callback adapter automatically intercepts LangChain execution signals and translates them to canonical AeroGraph `TraceEvent` types:
|
|
106
|
+
|
|
107
|
+
- **LLM/Chat Starts** (`on_llm_start`, `on_chat_model_start`) $\rightarrow$ `PromptEvent`
|
|
108
|
+
- **LLM Ends** (`on_llm_end`) $\rightarrow$ `ResponseEvent` (includes streaming telemetry completion metrics: TTFT and tokens/sec when tokens are streamed)
|
|
109
|
+
- **Tool Starts** (`on_tool_start`) $\rightarrow$ `ToolCallEvent`
|
|
110
|
+
- **Tool Ends** (`on_tool_end`) $\rightarrow$ `ToolResultEvent`
|
|
111
|
+
- **Retriever Runs** (`on_retriever_start`, `on_retriever_end`) $\rightarrow$ `RetrieverEvent` (captures source documents, queries, and metadata)
|
|
112
|
+
- **Errors** (`on_llm_error`, `on_tool_error`, `on_chain_error`) $\rightarrow$ `ErrorEvent`
|
|
113
|
+
- **Custom Events** (`on_custom_event`) $\rightarrow$ `StateSnapshotEvent` and `CheckpointEvent` for LangGraph nodes
|
|
114
|
+
|
|
115
|
+
## License
|
|
116
|
+
|
|
117
|
+
Apache-2.0
|
|
@@ -0,0 +1,88 @@
|
|
|
1
|
+
# aerograph-langchain
|
|
2
|
+
|
|
3
|
+
LangChain callback adapter for [AeroGraph](https://github.com/SGcpu/AeroGraph).
|
|
4
|
+
|
|
5
|
+
Automatically records LangChain chain, LLM, tool, and retriever events as AeroGraph trace events.
|
|
6
|
+
|
|
7
|
+
## Installation
|
|
8
|
+
|
|
9
|
+
When published, you can install the adapter via pip:
|
|
10
|
+
|
|
11
|
+
```bash
|
|
12
|
+
pip install aerograph-langchain
|
|
13
|
+
```
|
|
14
|
+
|
|
15
|
+
For local testing or development, install the packages in editable mode from the repository root:
|
|
16
|
+
|
|
17
|
+
```bash
|
|
18
|
+
# Install the core SDK first
|
|
19
|
+
pip install -e python/aerograph-sdk
|
|
20
|
+
|
|
21
|
+
# Install the LangChain adapter
|
|
22
|
+
pip install -e python/aerograph-langchain
|
|
23
|
+
```
|
|
24
|
+
|
|
25
|
+
## Usage
|
|
26
|
+
|
|
27
|
+
Integrating `AeroGraphCallbackHandler` into your existing LangChain codebase is simple. Setup a `FlightRecorder`, initialize the handler, and pass it to your chain or model invocations.
|
|
28
|
+
|
|
29
|
+
### Basic Example
|
|
30
|
+
|
|
31
|
+
```python
|
|
32
|
+
import asyncio
|
|
33
|
+
from langchain_core.messages import HumanMessage
|
|
34
|
+
from langchain_openai import ChatOpenAI
|
|
35
|
+
from aerograph_sdk.recorder import FlightRecorder
|
|
36
|
+
from aerograph_langchain.handler import AeroGraphCallbackHandler
|
|
37
|
+
|
|
38
|
+
async def main():
|
|
39
|
+
# 1. Initialize the FlightRecorder pointing to your collector
|
|
40
|
+
recorder = FlightRecorder(
|
|
41
|
+
endpoint="http://localhost:4317",
|
|
42
|
+
actor={"id": "my-agent", "name": "TravelPlanner"}
|
|
43
|
+
)
|
|
44
|
+
|
|
45
|
+
# 2. Initialize the AeroGraph callback handler
|
|
46
|
+
handler = AeroGraphCallbackHandler(recorder)
|
|
47
|
+
|
|
48
|
+
# 3. Create your LangChain chat model
|
|
49
|
+
model = ChatOpenAI(model="gpt-4o")
|
|
50
|
+
|
|
51
|
+
# 4. Invoke the model and pass the handler in callbacks
|
|
52
|
+
response = await model.ainvoke(
|
|
53
|
+
[HumanMessage(content="What are 3 fun things to do in San Francisco?")],
|
|
54
|
+
config={"callbacks": [handler]}
|
|
55
|
+
)
|
|
56
|
+
print(response.content)
|
|
57
|
+
|
|
58
|
+
if __name__ == "__main__":
|
|
59
|
+
asyncio.run(main())
|
|
60
|
+
```
|
|
61
|
+
|
|
62
|
+
### Advanced Usage with Chains and RAG
|
|
63
|
+
|
|
64
|
+
You can attach the callback handler at the chain execution level. LangChain automatically propagates the callbacks down to all sub-chains, LLMs, retrievers, and tool invocations.
|
|
65
|
+
|
|
66
|
+
```python
|
|
67
|
+
# Pass the handler to the chain invoke call
|
|
68
|
+
result = rag_chain.invoke(
|
|
69
|
+
"How do I configure the server?",
|
|
70
|
+
config={"callbacks": [handler]}
|
|
71
|
+
)
|
|
72
|
+
```
|
|
73
|
+
|
|
74
|
+
## Features and Event Mapping
|
|
75
|
+
|
|
76
|
+
The callback adapter automatically intercepts LangChain execution signals and translates them to canonical AeroGraph `TraceEvent` types:
|
|
77
|
+
|
|
78
|
+
- **LLM/Chat Starts** (`on_llm_start`, `on_chat_model_start`) $\rightarrow$ `PromptEvent`
|
|
79
|
+
- **LLM Ends** (`on_llm_end`) $\rightarrow$ `ResponseEvent` (includes streaming telemetry completion metrics: TTFT and tokens/sec when tokens are streamed)
|
|
80
|
+
- **Tool Starts** (`on_tool_start`) $\rightarrow$ `ToolCallEvent`
|
|
81
|
+
- **Tool Ends** (`on_tool_end`) $\rightarrow$ `ToolResultEvent`
|
|
82
|
+
- **Retriever Runs** (`on_retriever_start`, `on_retriever_end`) $\rightarrow$ `RetrieverEvent` (captures source documents, queries, and metadata)
|
|
83
|
+
- **Errors** (`on_llm_error`, `on_tool_error`, `on_chain_error`) $\rightarrow$ `ErrorEvent`
|
|
84
|
+
- **Custom Events** (`on_custom_event`) $\rightarrow$ `StateSnapshotEvent` and `CheckpointEvent` for LangGraph nodes
|
|
85
|
+
|
|
86
|
+
## License
|
|
87
|
+
|
|
88
|
+
Apache-2.0
|
|
@@ -0,0 +1,53 @@
|
|
|
1
|
+
import asyncio
|
|
2
|
+
from langchain_core.messages import HumanMessage
|
|
3
|
+
from langchain_core.language_models.chat_models import BaseChatModel
|
|
4
|
+
from langchain_core.outputs import ChatGeneration, ChatResult
|
|
5
|
+
|
|
6
|
+
from aerograph_sdk.recorder import FlightRecorder
|
|
7
|
+
from aerograph_langchain.handler import AeroGraphCallbackHandler
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
class MockChatModel(BaseChatModel):
|
|
11
|
+
"""A mock chat model for demonstration without an API key."""
|
|
12
|
+
|
|
13
|
+
@property
|
|
14
|
+
def _llm_type(self) -> str:
|
|
15
|
+
return "mock-chat"
|
|
16
|
+
|
|
17
|
+
def _generate(self, messages, stop=None, run_manager=None, **kwargs):
|
|
18
|
+
text = "This is a mock response to demonstrate AeroGraph tracing."
|
|
19
|
+
message = ChatGeneration(text=text, message={"content": text, "type": "ai"})
|
|
20
|
+
return ChatResult(generations=[message])
|
|
21
|
+
|
|
22
|
+
async def _agenerate(self, messages, stop=None, run_manager=None, **kwargs):
|
|
23
|
+
return self._generate(messages, stop, run_manager, **kwargs)
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
async def main():
|
|
27
|
+
# 1. Initialize the flight recorder
|
|
28
|
+
# By default, it connects to http://localhost:4000
|
|
29
|
+
recorder = FlightRecorder(api_url="http://localhost:4000")
|
|
30
|
+
|
|
31
|
+
# 2. Create the AeroGraph LangChain callback handler
|
|
32
|
+
aerograph_handler = AeroGraphCallbackHandler(recorder)
|
|
33
|
+
|
|
34
|
+
# 3. Create a model (mocked here, but typically ChatOpenAI etc)
|
|
35
|
+
model = MockChatModel()
|
|
36
|
+
|
|
37
|
+
print(f"Starting trace: {aerograph_handler.trace_id}")
|
|
38
|
+
|
|
39
|
+
# 4. Invoke with the callback
|
|
40
|
+
messages = [HumanMessage(content="Tell me about AeroGraph observability.")]
|
|
41
|
+
|
|
42
|
+
print("Invoking model...")
|
|
43
|
+
response = await model.ainvoke(messages, config={"callbacks": [aerograph_handler]})
|
|
44
|
+
|
|
45
|
+
print("Response received:", response)
|
|
46
|
+
|
|
47
|
+
# Wait a moment for background events to flush
|
|
48
|
+
await asyncio.sleep(0.5)
|
|
49
|
+
print("Done. Check the AeroGraph UI to view the trace.")
|
|
50
|
+
|
|
51
|
+
|
|
52
|
+
if __name__ == "__main__":
|
|
53
|
+
asyncio.run(main())
|
|
@@ -0,0 +1,56 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["hatchling"]
|
|
3
|
+
build-backend = "hatchling.build"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "aerograph-langchain"
|
|
7
|
+
version = "0.2.0"
|
|
8
|
+
description = "LangChain callback adapter for AeroGraph — automatically record LangChain traces."
|
|
9
|
+
readme = "README.md"
|
|
10
|
+
license = { text = "Apache-2.0" }
|
|
11
|
+
requires-python = ">=3.10"
|
|
12
|
+
keywords = ["ai", "agent", "tracing", "langchain", "aerograph"]
|
|
13
|
+
classifiers = [
|
|
14
|
+
"Development Status :: 3 - Alpha",
|
|
15
|
+
"Intended Audience :: Developers",
|
|
16
|
+
"License :: OSI Approved :: Apache Software License",
|
|
17
|
+
"Programming Language :: Python :: 3",
|
|
18
|
+
"Programming Language :: Python :: 3.10",
|
|
19
|
+
"Programming Language :: Python :: 3.11",
|
|
20
|
+
"Programming Language :: Python :: 3.12",
|
|
21
|
+
"Topic :: Software Development :: Libraries",
|
|
22
|
+
]
|
|
23
|
+
dependencies = [
|
|
24
|
+
"aerograph-sdk>=0.2.0",
|
|
25
|
+
"langchain-core>=0.2.0",
|
|
26
|
+
]
|
|
27
|
+
|
|
28
|
+
[project.optional-dependencies]
|
|
29
|
+
dev = [
|
|
30
|
+
"pytest>=8.0.0",
|
|
31
|
+
"pytest-asyncio>=0.24.0",
|
|
32
|
+
"langchain>=0.2.0",
|
|
33
|
+
"langchain-core>=0.2.0",
|
|
34
|
+
]
|
|
35
|
+
langgraph = [
|
|
36
|
+
"langgraph>=0.1.0",
|
|
37
|
+
]
|
|
38
|
+
|
|
39
|
+
[project.urls]
|
|
40
|
+
Homepage = "https://github.com/SGcpu/AeroGraph"
|
|
41
|
+
Repository = "https://github.com/SGcpu/AeroGraph"
|
|
42
|
+
Issues = "https://github.com/SGcpu/AeroGraph/issues"
|
|
43
|
+
|
|
44
|
+
[tool.hatch.build.targets.wheel]
|
|
45
|
+
packages = ["src/aerograph_langchain"]
|
|
46
|
+
|
|
47
|
+
[tool.pytest.ini_options]
|
|
48
|
+
testpaths = ["tests"]
|
|
49
|
+
asyncio_mode = "auto"
|
|
50
|
+
asyncio_default_fixture_loop_scope = "function"
|
|
51
|
+
markers = [
|
|
52
|
+
"integration: marks tests as integration tests (requires external services)",
|
|
53
|
+
]
|
|
54
|
+
|
|
55
|
+
[tool.uv.sources]
|
|
56
|
+
aerograph-sdk = { path = "../aerograph-sdk", editable = true }
|
|
@@ -0,0 +1,252 @@
|
|
|
1
|
+
import uuid
|
|
2
|
+
from typing import Any, Dict, List, Optional
|
|
3
|
+
from langchain_core.callbacks import BaseCallbackHandler
|
|
4
|
+
from langchain_core.outputs import LLMResult
|
|
5
|
+
from langchain_core.messages import BaseMessage
|
|
6
|
+
from langchain_core.documents import Document
|
|
7
|
+
|
|
8
|
+
from aerograph_sdk.recorder import FlightRecorder
|
|
9
|
+
from aerograph_sdk.ids import new_trace_id
|
|
10
|
+
|
|
11
|
+
from .mapping import map_llm_start, map_llm_end, map_tool_start, map_tool_end, map_error
|
|
12
|
+
from .streaming import StreamingTracker
|
|
13
|
+
from .retriever import RetrieverTracker
|
|
14
|
+
from .langgraph import map_state_snapshot, map_checkpoint
|
|
15
|
+
|
|
16
|
+
|
|
17
|
+
class AeroGraphCallbackHandler(BaseCallbackHandler):
|
|
18
|
+
"""Callback Handler that records LangChain runs to AeroGraph."""
|
|
19
|
+
|
|
20
|
+
def __init__(
|
|
21
|
+
self, recorder: FlightRecorder, trace_id: Optional[str] = None
|
|
22
|
+
) -> None:
|
|
23
|
+
super().__init__()
|
|
24
|
+
self.recorder = recorder
|
|
25
|
+
self.trace_id = trace_id or new_trace_id()
|
|
26
|
+
self.streaming_tracker = StreamingTracker()
|
|
27
|
+
self.retriever_tracker = RetrieverTracker()
|
|
28
|
+
|
|
29
|
+
def on_llm_start(
|
|
30
|
+
self,
|
|
31
|
+
serialized: Dict[str, Any],
|
|
32
|
+
prompts: List[str],
|
|
33
|
+
*,
|
|
34
|
+
run_id: uuid.UUID,
|
|
35
|
+
parent_run_id: Optional[uuid.UUID] = None,
|
|
36
|
+
tags: Optional[List[str]] = None,
|
|
37
|
+
metadata: Optional[Dict[str, Any]] = None,
|
|
38
|
+
**kwargs: Any,
|
|
39
|
+
) -> Any:
|
|
40
|
+
event = map_llm_start(
|
|
41
|
+
serialized=serialized,
|
|
42
|
+
prompts=prompts,
|
|
43
|
+
run_id=run_id,
|
|
44
|
+
parent_run_id=parent_run_id,
|
|
45
|
+
trace_id=self.trace_id,
|
|
46
|
+
)
|
|
47
|
+
self.recorder.emit(event)
|
|
48
|
+
self.streaming_tracker.on_llm_start(run_id)
|
|
49
|
+
|
|
50
|
+
def on_chat_model_start(
|
|
51
|
+
self,
|
|
52
|
+
serialized: Dict[str, Any],
|
|
53
|
+
messages: List[List[BaseMessage]],
|
|
54
|
+
*,
|
|
55
|
+
run_id: uuid.UUID,
|
|
56
|
+
parent_run_id: Optional[uuid.UUID] = None,
|
|
57
|
+
tags: Optional[List[str]] = None,
|
|
58
|
+
metadata: Optional[Dict[str, Any]] = None,
|
|
59
|
+
**kwargs: Any,
|
|
60
|
+
) -> Any:
|
|
61
|
+
event = map_llm_start(
|
|
62
|
+
serialized=serialized,
|
|
63
|
+
messages=messages,
|
|
64
|
+
run_id=run_id,
|
|
65
|
+
parent_run_id=parent_run_id,
|
|
66
|
+
trace_id=self.trace_id,
|
|
67
|
+
)
|
|
68
|
+
self.recorder.emit(event)
|
|
69
|
+
self.streaming_tracker.on_llm_start(run_id)
|
|
70
|
+
|
|
71
|
+
def on_llm_new_token(
|
|
72
|
+
self,
|
|
73
|
+
token: str,
|
|
74
|
+
*,
|
|
75
|
+
chunk: Optional[Any] = None,
|
|
76
|
+
run_id: uuid.UUID,
|
|
77
|
+
parent_run_id: Optional[uuid.UUID] = None,
|
|
78
|
+
**kwargs: Any,
|
|
79
|
+
) -> Any:
|
|
80
|
+
self.streaming_tracker.on_llm_new_token(run_id)
|
|
81
|
+
|
|
82
|
+
def on_llm_end(
|
|
83
|
+
self,
|
|
84
|
+
response: LLMResult,
|
|
85
|
+
*,
|
|
86
|
+
run_id: uuid.UUID,
|
|
87
|
+
parent_run_id: Optional[uuid.UUID] = None,
|
|
88
|
+
**kwargs: Any,
|
|
89
|
+
) -> Any:
|
|
90
|
+
event = map_llm_end(
|
|
91
|
+
serialized={
|
|
92
|
+
"name": "LLM"
|
|
93
|
+
}, # Best effort, LangChain doesn't pass serialized to end
|
|
94
|
+
response=response,
|
|
95
|
+
run_id=run_id,
|
|
96
|
+
parent_run_id=parent_run_id,
|
|
97
|
+
trace_id=self.trace_id,
|
|
98
|
+
)
|
|
99
|
+
telemetry = self.streaming_tracker.on_llm_end(run_id)
|
|
100
|
+
if telemetry:
|
|
101
|
+
# We recreate the payload with telemetry
|
|
102
|
+
event.payload.streamingTelemetry = telemetry
|
|
103
|
+
self.recorder.emit(event)
|
|
104
|
+
|
|
105
|
+
def on_tool_start(
|
|
106
|
+
self,
|
|
107
|
+
serialized: Dict[str, Any],
|
|
108
|
+
input_str: str,
|
|
109
|
+
*,
|
|
110
|
+
run_id: uuid.UUID,
|
|
111
|
+
parent_run_id: Optional[uuid.UUID] = None,
|
|
112
|
+
tags: Optional[List[str]] = None,
|
|
113
|
+
metadata: Optional[Dict[str, Any]] = None,
|
|
114
|
+
inputs: Optional[Dict[str, Any]] = None,
|
|
115
|
+
**kwargs: Any,
|
|
116
|
+
) -> Any:
|
|
117
|
+
event = map_tool_start(
|
|
118
|
+
serialized=serialized,
|
|
119
|
+
input_str=input_str,
|
|
120
|
+
run_id=run_id,
|
|
121
|
+
parent_run_id=parent_run_id,
|
|
122
|
+
trace_id=self.trace_id,
|
|
123
|
+
inputs=inputs,
|
|
124
|
+
)
|
|
125
|
+
self.recorder.emit(event)
|
|
126
|
+
|
|
127
|
+
def on_tool_end(
|
|
128
|
+
self,
|
|
129
|
+
output: Any,
|
|
130
|
+
*,
|
|
131
|
+
run_id: uuid.UUID,
|
|
132
|
+
parent_run_id: Optional[uuid.UUID] = None,
|
|
133
|
+
**kwargs: Any,
|
|
134
|
+
) -> Any:
|
|
135
|
+
event = map_tool_end(
|
|
136
|
+
serialized={"name": "Tool"}, # Best effort
|
|
137
|
+
output=output,
|
|
138
|
+
run_id=run_id,
|
|
139
|
+
parent_run_id=parent_run_id,
|
|
140
|
+
trace_id=self.trace_id,
|
|
141
|
+
)
|
|
142
|
+
self.recorder.emit(event)
|
|
143
|
+
|
|
144
|
+
def on_retriever_start(
|
|
145
|
+
self,
|
|
146
|
+
serialized: Dict[str, Any],
|
|
147
|
+
query: str,
|
|
148
|
+
*,
|
|
149
|
+
run_id: uuid.UUID,
|
|
150
|
+
parent_run_id: Optional[uuid.UUID] = None,
|
|
151
|
+
tags: Optional[List[str]] = None,
|
|
152
|
+
metadata: Optional[Dict[str, Any]] = None,
|
|
153
|
+
**kwargs: Any,
|
|
154
|
+
) -> Any:
|
|
155
|
+
self.retriever_tracker.on_retriever_start(run_id, query, serialized)
|
|
156
|
+
|
|
157
|
+
def on_retriever_end(
|
|
158
|
+
self,
|
|
159
|
+
documents: List[Document],
|
|
160
|
+
*,
|
|
161
|
+
run_id: uuid.UUID,
|
|
162
|
+
parent_run_id: Optional[uuid.UUID] = None,
|
|
163
|
+
**kwargs: Any,
|
|
164
|
+
) -> Any:
|
|
165
|
+
event = self.retriever_tracker.on_retriever_end(
|
|
166
|
+
run_id=run_id,
|
|
167
|
+
documents=documents,
|
|
168
|
+
trace_id=self.trace_id,
|
|
169
|
+
parent_run_id=parent_run_id,
|
|
170
|
+
)
|
|
171
|
+
if event:
|
|
172
|
+
self.recorder.emit(event)
|
|
173
|
+
|
|
174
|
+
def on_custom_event(
|
|
175
|
+
self,
|
|
176
|
+
name: str,
|
|
177
|
+
data: Any,
|
|
178
|
+
*,
|
|
179
|
+
run_id: uuid.UUID,
|
|
180
|
+
tags: Optional[List[str]] = None,
|
|
181
|
+
metadata: Optional[Dict[str, Any]] = None,
|
|
182
|
+
**kwargs: Any,
|
|
183
|
+
) -> Any:
|
|
184
|
+
"""Handle custom events, particularly LangGraph state emissions if mapped here."""
|
|
185
|
+
if name == "langgraph_state_snapshot":
|
|
186
|
+
event = map_state_snapshot(
|
|
187
|
+
run_id=run_id,
|
|
188
|
+
trace_id=self.trace_id,
|
|
189
|
+
node_name=data.get("node_name", "unknown"),
|
|
190
|
+
state_hash=data.get("state_hash", ""),
|
|
191
|
+
state_diff=data.get("state_diff", {}),
|
|
192
|
+
full_state=data.get("full_state", {}),
|
|
193
|
+
removed_keys=data.get("removed_keys", []),
|
|
194
|
+
)
|
|
195
|
+
self.recorder.emit(event)
|
|
196
|
+
elif name == "langgraph_checkpoint":
|
|
197
|
+
event = map_checkpoint(
|
|
198
|
+
run_id=run_id,
|
|
199
|
+
trace_id=self.trace_id,
|
|
200
|
+
checkpoint_id=data.get("checkpoint_id", ""),
|
|
201
|
+
reason=data.get("reason", ""),
|
|
202
|
+
state=data.get("state", {}),
|
|
203
|
+
)
|
|
204
|
+
self.recorder.emit(event)
|
|
205
|
+
|
|
206
|
+
def on_llm_error(
|
|
207
|
+
self,
|
|
208
|
+
error: BaseException,
|
|
209
|
+
*,
|
|
210
|
+
run_id: uuid.UUID,
|
|
211
|
+
parent_run_id: Optional[uuid.UUID] = None,
|
|
212
|
+
**kwargs: Any,
|
|
213
|
+
) -> Any:
|
|
214
|
+
event = map_error(
|
|
215
|
+
error=error,
|
|
216
|
+
run_id=run_id,
|
|
217
|
+
trace_id=self.trace_id,
|
|
218
|
+
parent_run_id=parent_run_id
|
|
219
|
+
)
|
|
220
|
+
self.recorder.emit(event)
|
|
221
|
+
|
|
222
|
+
def on_tool_error(
|
|
223
|
+
self,
|
|
224
|
+
error: BaseException,
|
|
225
|
+
*,
|
|
226
|
+
run_id: uuid.UUID,
|
|
227
|
+
parent_run_id: Optional[uuid.UUID] = None,
|
|
228
|
+
**kwargs: Any,
|
|
229
|
+
) -> Any:
|
|
230
|
+
event = map_error(
|
|
231
|
+
error=error,
|
|
232
|
+
run_id=run_id,
|
|
233
|
+
trace_id=self.trace_id,
|
|
234
|
+
parent_run_id=parent_run_id
|
|
235
|
+
)
|
|
236
|
+
self.recorder.emit(event)
|
|
237
|
+
|
|
238
|
+
def on_chain_error(
|
|
239
|
+
self,
|
|
240
|
+
error: BaseException,
|
|
241
|
+
*,
|
|
242
|
+
run_id: uuid.UUID,
|
|
243
|
+
parent_run_id: Optional[uuid.UUID] = None,
|
|
244
|
+
**kwargs: Any,
|
|
245
|
+
) -> Any:
|
|
246
|
+
event = map_error(
|
|
247
|
+
error=error,
|
|
248
|
+
run_id=run_id,
|
|
249
|
+
trace_id=self.trace_id,
|
|
250
|
+
parent_run_id=parent_run_id
|
|
251
|
+
)
|
|
252
|
+
self.recorder.emit(event)
|
|
@@ -0,0 +1,56 @@
|
|
|
1
|
+
import uuid
|
|
2
|
+
from typing import Any, Dict, List, Optional
|
|
3
|
+
|
|
4
|
+
from aerograph_sdk.events import build_state_snapshot_event, build_checkpoint_event
|
|
5
|
+
from aerograph_sdk.contracts.generated import StateSnapshotEvent, CheckpointEvent
|
|
6
|
+
from aerograph_langchain.span_ids import derive_span_id
|
|
7
|
+
|
|
8
|
+
|
|
9
|
+
def map_state_snapshot(
|
|
10
|
+
run_id: uuid.UUID,
|
|
11
|
+
trace_id: str,
|
|
12
|
+
node_name: str,
|
|
13
|
+
state_hash: str,
|
|
14
|
+
state_diff: Dict[str, Any],
|
|
15
|
+
full_state: Dict[str, Any],
|
|
16
|
+
removed_keys: Optional[List[str]] = None,
|
|
17
|
+
parent_run_id: Optional[uuid.UUID] = None,
|
|
18
|
+
) -> StateSnapshotEvent:
|
|
19
|
+
span_id = derive_span_id(run_id)
|
|
20
|
+
parent_span_id = derive_span_id(parent_run_id) if parent_run_id else None
|
|
21
|
+
|
|
22
|
+
return build_state_snapshot_event(
|
|
23
|
+
trace_id=trace_id,
|
|
24
|
+
span_id=span_id,
|
|
25
|
+
parent_span_id=parent_span_id,
|
|
26
|
+
actor_id="langgraph",
|
|
27
|
+
actor_name="LangGraph",
|
|
28
|
+
node_name=node_name,
|
|
29
|
+
state_hash=state_hash,
|
|
30
|
+
state_diff=state_diff,
|
|
31
|
+
full_state=full_state,
|
|
32
|
+
removed_keys=removed_keys or [],
|
|
33
|
+
)
|
|
34
|
+
|
|
35
|
+
|
|
36
|
+
def map_checkpoint(
|
|
37
|
+
run_id: uuid.UUID,
|
|
38
|
+
trace_id: str,
|
|
39
|
+
checkpoint_id: str,
|
|
40
|
+
reason: str,
|
|
41
|
+
state: Dict[str, Any],
|
|
42
|
+
parent_run_id: Optional[uuid.UUID] = None,
|
|
43
|
+
) -> CheckpointEvent:
|
|
44
|
+
span_id = derive_span_id(run_id)
|
|
45
|
+
parent_span_id = derive_span_id(parent_run_id) if parent_run_id else None
|
|
46
|
+
|
|
47
|
+
return build_checkpoint_event(
|
|
48
|
+
trace_id=trace_id,
|
|
49
|
+
span_id=span_id,
|
|
50
|
+
parent_span_id=parent_span_id,
|
|
51
|
+
actor_id="langgraph",
|
|
52
|
+
actor_name="LangGraph",
|
|
53
|
+
checkpoint_id=checkpoint_id,
|
|
54
|
+
reason=reason,
|
|
55
|
+
state=state,
|
|
56
|
+
)
|