agentx-python 0.4.11__tar.gz → 0.4.12__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.
Files changed (46) hide show
  1. {agentx_python-0.4.11/agentx_python.egg-info → agentx_python-0.4.12}/PKG-INFO +15 -1
  2. agentx_python-0.4.12/agentx/__init__.py +31 -0
  3. {agentx_python-0.4.11 → agentx_python-0.4.12}/agentx/agentx.py +9 -0
  4. agentx_python-0.4.12/agentx/exceptions.py +46 -0
  5. agentx_python-0.4.12/agentx/integrations/__init__.py +7 -0
  6. agentx_python-0.4.12/agentx/integrations/anthropic.py +152 -0
  7. agentx_python-0.4.12/agentx/integrations/crewai.py +106 -0
  8. agentx_python-0.4.12/agentx/integrations/langchain.py +249 -0
  9. agentx_python-0.4.12/agentx/integrations/openai_agents.py +125 -0
  10. {agentx_python-0.4.11 → agentx_python-0.4.12}/agentx/resources/conversation.py +4 -4
  11. agentx_python-0.4.12/agentx/tracing/__init__.py +21 -0
  12. agentx_python-0.4.12/agentx/tracing/ci_types.py +66 -0
  13. agentx_python-0.4.12/agentx/tracing/ingest_client.py +312 -0
  14. agentx_python-0.4.12/agentx/tracing/tracer.py +465 -0
  15. agentx_python-0.4.12/agentx/version.py +1 -0
  16. {agentx_python-0.4.11 → agentx_python-0.4.12/agentx_python.egg-info}/PKG-INFO +15 -1
  17. {agentx_python-0.4.11 → agentx_python-0.4.12}/agentx_python.egg-info/SOURCES.txt +10 -0
  18. agentx_python-0.4.12/agentx_python.egg-info/requires.txt +23 -0
  19. {agentx_python-0.4.11 → agentx_python-0.4.12}/setup.py +12 -0
  20. agentx_python-0.4.11/agentx/__init__.py +0 -13
  21. agentx_python-0.4.11/agentx/version.py +0 -1
  22. agentx_python-0.4.11/agentx_python.egg-info/requires.txt +0 -5
  23. {agentx_python-0.4.11 → agentx_python-0.4.12}/LICENSE +0 -0
  24. {agentx_python-0.4.11 → agentx_python-0.4.12}/README.md +0 -0
  25. {agentx_python-0.4.11 → agentx_python-0.4.12}/agentx/evaluations/__init__.py +0 -0
  26. {agentx_python-0.4.11 → agentx_python-0.4.12}/agentx/evaluations/_term.py +0 -0
  27. {agentx_python-0.4.11 → agentx_python-0.4.12}/agentx/evaluations/adapters/__init__.py +0 -0
  28. {agentx_python-0.4.11 → agentx_python-0.4.12}/agentx/evaluations/adapters/http_endpoint.py +0 -0
  29. {agentx_python-0.4.11 → agentx_python-0.4.12}/agentx/evaluations/adapters/precomputed.py +0 -0
  30. {agentx_python-0.4.11 → agentx_python-0.4.12}/agentx/evaluations/adapters/raw.py +0 -0
  31. {agentx_python-0.4.11 → agentx_python-0.4.12}/agentx/evaluations/client.py +0 -0
  32. {agentx_python-0.4.11 → agentx_python-0.4.12}/agentx/evaluations/datasets.py +0 -0
  33. {agentx_python-0.4.11 → agentx_python-0.4.12}/agentx/evaluations/models.py +0 -0
  34. {agentx_python-0.4.11 → agentx_python-0.4.12}/agentx/evaluations/redaction.py +0 -0
  35. {agentx_python-0.4.11 → agentx_python-0.4.12}/agentx/evaluations/reporting.py +0 -0
  36. {agentx_python-0.4.11 → agentx_python-0.4.12}/agentx/evaluations/results.py +0 -0
  37. {agentx_python-0.4.11 → agentx_python-0.4.12}/agentx/evaluations/runner.py +0 -0
  38. {agentx_python-0.4.11 → agentx_python-0.4.12}/agentx/evaluations/tracing.py +0 -0
  39. {agentx_python-0.4.11 → agentx_python-0.4.12}/agentx/resources/__init__.py +0 -0
  40. {agentx_python-0.4.11 → agentx_python-0.4.12}/agentx/resources/agent.py +0 -0
  41. {agentx_python-0.4.11 → agentx_python-0.4.12}/agentx/resources/workforce.py +0 -0
  42. {agentx_python-0.4.11 → agentx_python-0.4.12}/agentx/util.py +0 -0
  43. {agentx_python-0.4.11 → agentx_python-0.4.12}/agentx_python.egg-info/dependency_links.txt +0 -0
  44. {agentx_python-0.4.11 → agentx_python-0.4.12}/agentx_python.egg-info/top_level.txt +0 -0
  45. {agentx_python-0.4.11 → agentx_python-0.4.12}/setup.cfg +0 -0
  46. {agentx_python-0.4.11 → agentx_python-0.4.12}/tests/test_integration.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: agentx-python
3
- Version: 0.4.11
3
+ Version: 0.4.12
4
4
  Summary: Official Python SDK for AgentX (https://www.agentx.so/)
5
5
  Home-page: https://github.com/AgentX-ai/AgentX-python
6
6
  Author: Robin Wang and AgentX Team
@@ -16,6 +16,19 @@ Requires-Dist: certifi
16
16
  Requires-Dist: requests
17
17
  Requires-Dist: pydantic
18
18
  Requires-Dist: pydantic_core
19
+ Provides-Extra: langchain
20
+ Requires-Dist: langchain-core>=0.1.0; extra == "langchain"
21
+ Provides-Extra: crewai
22
+ Requires-Dist: crewai>=0.80.0; extra == "crewai"
23
+ Provides-Extra: openai-agents
24
+ Requires-Dist: openai-agents>=0.0.3; extra == "openai-agents"
25
+ Provides-Extra: anthropic
26
+ Requires-Dist: anthropic>=0.25.0; extra == "anthropic"
27
+ Provides-Extra: all
28
+ Requires-Dist: langchain-core>=0.1.0; extra == "all"
29
+ Requires-Dist: crewai>=0.80.0; extra == "all"
30
+ Requires-Dist: openai-agents>=0.0.3; extra == "all"
31
+ Requires-Dist: anthropic>=0.25.0; extra == "all"
19
32
  Dynamic: author
20
33
  Dynamic: author-email
21
34
  Dynamic: classifier
@@ -23,6 +36,7 @@ Dynamic: description
23
36
  Dynamic: description-content-type
24
37
  Dynamic: home-page
25
38
  Dynamic: license-file
39
+ Dynamic: provides-extra
26
40
  Dynamic: requires-dist
27
41
  Dynamic: requires-python
28
42
  Dynamic: summary
@@ -0,0 +1,31 @@
1
+ import logging
2
+
3
+ from agentx.agentx import AgentX
4
+ from agentx.version import VERSION
5
+ from agentx.exceptions import (
6
+ AgentXError,
7
+ AgentXAuthError,
8
+ AgentXAPIError,
9
+ DatasetNotFound,
10
+ CINotEnabled,
11
+ CIRunExpired,
12
+ CIGateFailure,
13
+ )
14
+
15
+ logging.basicConfig(
16
+ level=logging.INFO,
17
+ format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
18
+ datefmt="%Y-%m-%d %H:%M:%S %Z",
19
+ )
20
+
21
+ __all__ = [
22
+ "AgentX",
23
+ "AgentXError",
24
+ "AgentXAuthError",
25
+ "AgentXAPIError",
26
+ "DatasetNotFound",
27
+ "CINotEnabled",
28
+ "CIRunExpired",
29
+ "CIGateFailure",
30
+ ]
31
+ __version__ = VERSION
@@ -22,6 +22,8 @@ class AgentX:
22
22
 
23
23
  from agentx.evaluations.client import EvaluationsClient
24
24
  from agentx.evaluations.runner import EvaluationsRunner
25
+ from agentx.tracing.ingest_client import IngestClient
26
+ from agentx.tracing.tracer import Tracer
25
27
  from agentx.version import VERSION
26
28
 
27
29
  _eval_client = EvaluationsClient(
@@ -31,6 +33,13 @@ class AgentX:
31
33
  )
32
34
  self.evaluations = EvaluationsRunner(_eval_client)
33
35
 
36
+ _ingest_client = IngestClient(
37
+ api_key=self.api_key,
38
+ sdk_version=VERSION,
39
+ base_url=self.base_url,
40
+ )
41
+ self.tracer = Tracer(_ingest_client)
42
+
34
43
  @classmethod
35
44
  def from_env(cls) -> "AgentX":
36
45
  """Create an AgentX client using AGENTX_API_KEY (and optionally AGENTX_API_BASE_URL) from the environment."""
@@ -0,0 +1,46 @@
1
+ """AgentX SDK exceptions."""
2
+ from __future__ import annotations
3
+
4
+ from typing import TYPE_CHECKING
5
+
6
+ if TYPE_CHECKING:
7
+ from agentx.tracing.ci_types import CIRunResult
8
+
9
+
10
+ class AgentXError(Exception):
11
+ """Base class for all AgentX SDK errors."""
12
+
13
+
14
+ class AgentXAuthError(AgentXError):
15
+ """Invalid or missing API key."""
16
+
17
+
18
+ class AgentXAPIError(AgentXError):
19
+ """Unexpected API error."""
20
+
21
+ def __init__(self, message: str, status_code: int | None = None) -> None:
22
+ super().__init__(message)
23
+ self.status_code = status_code
24
+
25
+
26
+ class DatasetNotFound(AgentXError):
27
+ """Dataset ID does not exist or is not accessible from this API key."""
28
+
29
+
30
+ class CINotEnabled(AgentXError):
31
+ """Dataset exists but ci.enabled is false. Enable CI in the dataset settings."""
32
+
33
+
34
+ class CIRunExpired(AgentXError):
35
+ """CI run was not finalized within the 2-hour window."""
36
+
37
+
38
+ class CIGateFailure(AgentXError):
39
+ """Gate result is 'fail'. Raised when fail_on_gate=True."""
40
+
41
+ def __init__(self, result: "CIRunResult") -> None:
42
+ super().__init__(
43
+ f"CI gate failed: {result.pass_rate:.0%} passed "
44
+ f"({result.passed_questions}/{result.total_questions} questions)"
45
+ )
46
+ self.result = result
@@ -0,0 +1,7 @@
1
+ # Framework-specific integrations for production tracing.
2
+ # Each sub-module is lazily importable so the core SDK has no extra dependencies.
3
+ #
4
+ # from agentx.integrations.langchain import AgentXCallbackHandler
5
+ # from agentx.integrations.crewai import AgentXCrewObserver
6
+ # from agentx.integrations.openai_agents import AgentXTracingProcessor
7
+ # from agentx.integrations.anthropic import patch_anthropic_client
@@ -0,0 +1,152 @@
1
+ """
2
+ Anthropic SDK integration for AgentX production tracing.
3
+
4
+ Usage::
5
+
6
+ from agentx.integrations.anthropic import patch_anthropic_client
7
+ import anthropic
8
+
9
+ client = anthropic.Anthropic()
10
+ patch_anthropic_client(client, agentx.tracer, name="claude-agent")
11
+
12
+ # All subsequent client.messages.create() calls are now traced automatically.
13
+
14
+ Requires: ``pip install agentx[anthropic]``
15
+ """
16
+ from __future__ import annotations
17
+
18
+ import time
19
+ from typing import Any, Dict, Optional
20
+
21
+ from agentx.tracing.tracer import Tracer, _safe_serialize
22
+
23
+
24
+ def patch_anthropic_client(
25
+ client: Any,
26
+ tracer: Tracer,
27
+ name: str = "anthropic-agent",
28
+ metadata: Optional[Dict[str, Any]] = None,
29
+ session_id: Optional[str] = None,
30
+ ) -> None:
31
+ """
32
+ Monkey-patch ``client.messages.create`` and ``client.messages.stream``
33
+ (if present) to automatically send a trace for every call.
34
+
35
+ The original method is still called and its return value is passed through
36
+ unchanged so nothing in the caller needs to change.
37
+ """
38
+ messages = getattr(client, "messages", None)
39
+ if messages is None:
40
+ raise ValueError("Provided client does not have a .messages attribute")
41
+
42
+ _patch_create(messages, tracer, name, metadata, session_id)
43
+
44
+ # stream is optional (not present in all versions)
45
+ if hasattr(messages, "stream"):
46
+ _patch_stream(messages, tracer, name, metadata, session_id)
47
+
48
+
49
+ def _patch_create(
50
+ messages_resource: Any,
51
+ tracer: Tracer,
52
+ name: str,
53
+ metadata: Optional[Dict[str, Any]],
54
+ session_id: Optional[str],
55
+ ) -> None:
56
+ original = messages_resource.create
57
+ if getattr(original, "_agentx_patched", False):
58
+ return # already patched
59
+
60
+ def patched_create(*args, **kwargs):
61
+ start = time.time()
62
+ error: Optional[str] = None
63
+ response = None
64
+ try:
65
+ response = original(*args, **kwargs)
66
+ return response
67
+ except Exception as exc:
68
+ error = str(exc)
69
+ raise
70
+ finally:
71
+ latency_ms = int((time.time() - start) * 1000)
72
+ input_messages = kwargs.get("messages") or (args[0] if args else None)
73
+ model = kwargs.get("model")
74
+ output = None
75
+ if response is not None:
76
+ try:
77
+ output = response.content[0].text if response.content else None
78
+ except Exception:
79
+ output = str(response)[:500]
80
+ tracer._send(
81
+ name=name,
82
+ input=_safe_serialize(input_messages),
83
+ output=output,
84
+ latency_ms=latency_ms,
85
+ error=error,
86
+ framework="anthropic",
87
+ model=model,
88
+ metadata=metadata,
89
+ session_id=session_id,
90
+ )
91
+
92
+ patched_create._agentx_patched = True
93
+ messages_resource.create = patched_create
94
+
95
+
96
+ def _patch_stream(
97
+ messages_resource: Any,
98
+ tracer: Tracer,
99
+ name: str,
100
+ metadata: Optional[Dict[str, Any]],
101
+ session_id: Optional[str],
102
+ ) -> None:
103
+ original_stream = messages_resource.stream
104
+ if getattr(original_stream, "_agentx_patched", False):
105
+ return
106
+
107
+ def patched_stream(*args, **kwargs):
108
+ start = time.time()
109
+ ctx = original_stream(*args, **kwargs)
110
+
111
+ class _TracedStream:
112
+ """Thin wrapper that records timing when the stream context exits."""
113
+
114
+ def __enter__(self_inner):
115
+ return ctx.__enter__()
116
+
117
+ def __exit__(self_inner, exc_type, exc_val, tb):
118
+ result = ctx.__exit__(exc_type, exc_val, tb)
119
+ latency_ms = int((time.time() - start) * 1000)
120
+ error = str(exc_val) if exc_val else None
121
+ output = None
122
+ try:
123
+ final = ctx.get_final_message()
124
+ output = final.content[0].text if final.content else None
125
+ except Exception:
126
+ pass
127
+ tracer._send(
128
+ name=name,
129
+ input=_safe_serialize(kwargs.get("messages")),
130
+ output=output,
131
+ latency_ms=latency_ms,
132
+ error=error,
133
+ framework="anthropic",
134
+ model=kwargs.get("model"),
135
+ metadata=metadata,
136
+ session_id=session_id,
137
+ )
138
+ return result
139
+
140
+ def __iter__(self_inner):
141
+ return iter(ctx)
142
+
143
+ def __aiter__(self_inner):
144
+ return aiter(ctx)
145
+
146
+ def __getattr__(self_inner, item):
147
+ return getattr(ctx, item)
148
+
149
+ return _TracedStream()
150
+
151
+ patched_stream._agentx_patched = True
152
+ messages_resource.stream = patched_stream
@@ -0,0 +1,106 @@
1
+ """
2
+ CrewAI integration for AgentX production tracing.
3
+
4
+ Usage::
5
+
6
+ from agentx.integrations.crewai import AgentXCrewObserver
7
+
8
+ observer = AgentXCrewObserver(agentx.tracer, name="my-crew")
9
+ result = observer.kickoff(crew, inputs={"topic": "AI"})
10
+
11
+ Or as a context manager around your own kickoff::
12
+
13
+ with observer.observe(name="my-crew", input={"topic": "AI"}) as span:
14
+ result = crew.kickoff(inputs={"topic": "AI"})
15
+ span.output = result.raw
16
+
17
+ Requires: ``pip install agentx[crewai]``
18
+ """
19
+ from __future__ import annotations
20
+
21
+ import time
22
+ from typing import Any, Dict, Optional
23
+
24
+ from agentx.tracing.tracer import Tracer, _safe_serialize
25
+
26
+
27
+ class AgentXCrewObserver:
28
+ """
29
+ Wraps a CrewAI ``Crew.kickoff()`` call and sends one trace per execution.
30
+ Does not require any CrewAI version-specific event system.
31
+ """
32
+
33
+ def __init__(
34
+ self,
35
+ tracer: Tracer,
36
+ name: str = "crewai-crew",
37
+ metadata: Optional[Dict[str, Any]] = None,
38
+ session_id: Optional[str] = None,
39
+ ) -> None:
40
+ self._tracer = tracer
41
+ self._name = name
42
+ self._metadata = metadata
43
+ self._session_id = session_id
44
+
45
+ def kickoff(self, crew: Any, inputs: Optional[Dict[str, Any]] = None) -> Any:
46
+ """
47
+ Call ``crew.kickoff(inputs=inputs)``, capture result, and send a trace.
48
+ Returns the raw CrewAI ``CrewOutput`` object unchanged.
49
+ """
50
+ start = time.time()
51
+ error: Optional[str] = None
52
+ result = None
53
+ try:
54
+ result = crew.kickoff(inputs=inputs or {})
55
+ return result
56
+ except Exception as exc:
57
+ error = str(exc)
58
+ raise
59
+ finally:
60
+ latency_ms = int((time.time() - start) * 1000)
61
+ output = None
62
+ if result is not None:
63
+ output = getattr(result, "raw", None) or _safe_serialize(result)
64
+
65
+ # Collect task outputs as tool_calls for observability
66
+ tool_calls = []
67
+ if result is not None:
68
+ task_outputs = getattr(result, "tasks_output", []) or []
69
+ for task_out in task_outputs:
70
+ tool_calls.append(
71
+ {
72
+ "name": getattr(task_out, "description", "task")[:100],
73
+ "output": str(getattr(task_out, "raw", ""))[:500],
74
+ }
75
+ )
76
+
77
+ self._tracer._send(
78
+ name=self._name,
79
+ input=_safe_serialize(inputs) if inputs else None,
80
+ output=output,
81
+ latency_ms=latency_ms,
82
+ error=error,
83
+ framework="crewai",
84
+ tool_calls=tool_calls or None,
85
+ metadata=self._metadata,
86
+ session_id=self._session_id,
87
+ )
88
+
89
+ def observe(
90
+ self,
91
+ name: Optional[str] = None,
92
+ input: Optional[Any] = None,
93
+ metadata: Optional[Dict[str, Any]] = None,
94
+ session_id: Optional[str] = None,
95
+ ):
96
+ """Return a context-manager span for manual kickoff wrapping."""
97
+ from agentx.tracing.tracer import _TraceSpan
98
+
99
+ return _TraceSpan(
100
+ tracer=self._tracer,
101
+ name=name or self._name,
102
+ input=_safe_serialize(input) if input is not None else None,
103
+ metadata=metadata or self._metadata,
104
+ framework="crewai",
105
+ session_id=session_id or self._session_id,
106
+ )
@@ -0,0 +1,249 @@
1
+ """
2
+ LangChain integration for AgentX production tracing.
3
+
4
+ Usage::
5
+
6
+ from agentx.integrations.langchain import AgentXCallbackHandler
7
+
8
+ handler = AgentXCallbackHandler(agentx.tracer, name="my-chain")
9
+
10
+ # LCEL chain
11
+ chain.invoke({"query": q}, config={"callbacks": [handler]})
12
+
13
+ # AgentExecutor
14
+ agent.invoke({"input": q}, config={"callbacks": [handler]})
15
+
16
+ Requires: ``pip install agentx[langchain]``
17
+ """
18
+ from __future__ import annotations
19
+
20
+ import time
21
+ from typing import Any, Dict, List, Optional, Union
22
+ from uuid import UUID
23
+
24
+ from agentx.tracing.tracer import Tracer, _safe_serialize
25
+
26
+ try:
27
+ from langchain_core.callbacks.base import BaseCallbackHandler
28
+ from langchain_core.outputs import LLMResult
29
+ except ImportError as exc: # pragma: no cover
30
+ raise ImportError(
31
+ "langchain-core is required for AgentXCallbackHandler. "
32
+ "Install it with: pip install agentx[langchain]"
33
+ ) from exc
34
+
35
+
36
+ class AgentXCallbackHandler(BaseCallbackHandler):
37
+ """
38
+ LangChain callback handler that captures the top-level chain run and all
39
+ nested tool calls, then sends one trace per top-level chain invocation.
40
+ """
41
+
42
+ def __init__(
43
+ self,
44
+ tracer: Tracer,
45
+ name: str = "langchain-agent",
46
+ metadata: Optional[Dict[str, Any]] = None,
47
+ session_id: Optional[str] = None,
48
+ ) -> None:
49
+ super().__init__()
50
+ self._tracer = tracer
51
+ self._name = name
52
+ self._metadata = metadata
53
+ self._session_id = session_id
54
+
55
+ # Keyed by run_id (UUID) → state dict
56
+ self._runs: Dict[UUID, Dict[str, Any]] = {}
57
+ # Track which run_ids are top-level (no parent)
58
+ self._top_level: Dict[UUID, bool] = {}
59
+
60
+ # ------------------------------------------------------------------
61
+ # Chain lifecycle
62
+ # ------------------------------------------------------------------
63
+
64
+ def on_chain_start(
65
+ self,
66
+ serialized: Dict[str, Any],
67
+ inputs: Dict[str, Any],
68
+ *,
69
+ run_id: UUID,
70
+ parent_run_id: Optional[UUID] = None,
71
+ **kwargs,
72
+ ) -> None:
73
+ is_top = parent_run_id is None
74
+ self._top_level[run_id] = is_top
75
+ if is_top:
76
+ self._runs[run_id] = {
77
+ "start": time.time(),
78
+ "input": _safe_serialize(inputs),
79
+ "tool_calls": [],
80
+ "model": None,
81
+ }
82
+
83
+ def on_chain_end(
84
+ self,
85
+ outputs: Dict[str, Any],
86
+ *,
87
+ run_id: UUID,
88
+ parent_run_id: Optional[UUID] = None,
89
+ **kwargs,
90
+ ) -> None:
91
+ if not self._top_level.get(run_id):
92
+ return
93
+ state = self._runs.pop(run_id, None)
94
+ if state is None:
95
+ return
96
+ latency_ms = int((time.time() - state["start"]) * 1000)
97
+ self._tracer._send(
98
+ name=self._name,
99
+ input=state["input"],
100
+ output=_safe_serialize(outputs),
101
+ latency_ms=latency_ms,
102
+ framework="langchain",
103
+ model=state.get("model"),
104
+ tool_calls=state["tool_calls"] or None,
105
+ metadata=self._metadata,
106
+ session_id=self._session_id,
107
+ )
108
+ self._top_level.pop(run_id, None)
109
+
110
+ def on_chain_error(
111
+ self,
112
+ error: Union[Exception, KeyboardInterrupt],
113
+ *,
114
+ run_id: UUID,
115
+ parent_run_id: Optional[UUID] = None,
116
+ **kwargs,
117
+ ) -> None:
118
+ if not self._top_level.get(run_id):
119
+ return
120
+ state = self._runs.pop(run_id, None)
121
+ if state is None:
122
+ return
123
+ latency_ms = int((time.time() - state["start"]) * 1000)
124
+ self._tracer._send(
125
+ name=self._name,
126
+ input=state["input"],
127
+ error=str(error),
128
+ latency_ms=latency_ms,
129
+ framework="langchain",
130
+ model=state.get("model"),
131
+ tool_calls=state["tool_calls"] or None,
132
+ metadata=self._metadata,
133
+ session_id=self._session_id,
134
+ )
135
+ self._top_level.pop(run_id, None)
136
+
137
+ # ------------------------------------------------------------------
138
+ # LLM lifecycle (captures model name and token counts)
139
+ # ------------------------------------------------------------------
140
+
141
+ def on_llm_start(
142
+ self,
143
+ serialized: Dict[str, Any],
144
+ prompts: List[str],
145
+ *,
146
+ run_id: UUID,
147
+ parent_run_id: Optional[UUID] = None,
148
+ **kwargs,
149
+ ) -> None:
150
+ # Record LLM start time keyed by its own run_id
151
+ self._runs.setdefault(run_id, {})["llm_start"] = time.time()
152
+
153
+ # Propagate model name to the top-level run
154
+ top = self._find_top_ancestor(parent_run_id)
155
+ if top and not self._runs[top].get("model"):
156
+ model = (
157
+ serialized.get("kwargs", {}).get("model_name")
158
+ or serialized.get("kwargs", {}).get("model")
159
+ or serialized.get("id", [None])[-1]
160
+ )
161
+ if model:
162
+ self._runs[top]["model"] = str(model)
163
+
164
+ def on_llm_end(
165
+ self,
166
+ response: LLMResult,
167
+ *,
168
+ run_id: UUID,
169
+ parent_run_id: Optional[UUID] = None,
170
+ **kwargs,
171
+ ) -> None:
172
+ self._runs.pop(run_id, None)
173
+
174
+ # ------------------------------------------------------------------
175
+ # Tool lifecycle
176
+ # ------------------------------------------------------------------
177
+
178
+ def on_tool_start(
179
+ self,
180
+ serialized: Dict[str, Any],
181
+ input_str: str,
182
+ *,
183
+ run_id: UUID,
184
+ parent_run_id: Optional[UUID] = None,
185
+ **kwargs,
186
+ ) -> None:
187
+ self._runs[run_id] = {
188
+ "tool_name": serialized.get("name", "unknown"),
189
+ "tool_input": input_str,
190
+ "start": time.time(),
191
+ }
192
+
193
+ def on_tool_end(
194
+ self,
195
+ output: Any,
196
+ *,
197
+ run_id: UUID,
198
+ parent_run_id: Optional[UUID] = None,
199
+ **kwargs,
200
+ ) -> None:
201
+ state = self._runs.pop(run_id, None)
202
+ if state is None:
203
+ return
204
+ latency_ms = int((time.time() - state["start"]) * 1000)
205
+ tool_call = {
206
+ "name": state["tool_name"],
207
+ "input": state["tool_input"],
208
+ "output": str(output)[:500],
209
+ "latency_ms": latency_ms,
210
+ }
211
+ top = self._find_top_ancestor(parent_run_id)
212
+ if top and top in self._runs:
213
+ self._runs[top]["tool_calls"].append(tool_call)
214
+
215
+ def on_tool_error(
216
+ self,
217
+ error: Union[Exception, KeyboardInterrupt],
218
+ *,
219
+ run_id: UUID,
220
+ parent_run_id: Optional[UUID] = None,
221
+ **kwargs,
222
+ ) -> None:
223
+ state = self._runs.pop(run_id, None)
224
+ if state is None:
225
+ return
226
+ tool_call = {
227
+ "name": state.get("tool_name", "unknown"),
228
+ "input": state.get("tool_input"),
229
+ "output": f"ERROR: {error}",
230
+ "latency_ms": int((time.time() - state["start"]) * 1000),
231
+ }
232
+ top = self._find_top_ancestor(parent_run_id)
233
+ if top and top in self._runs:
234
+ self._runs[top]["tool_calls"].append(tool_call)
235
+
236
+ # ------------------------------------------------------------------
237
+ # Helpers
238
+ # ------------------------------------------------------------------
239
+
240
+ def _find_top_ancestor(self, parent_run_id: Optional[UUID]) -> Optional[UUID]:
241
+ """Walk up parent chain to find the top-level run_id."""
242
+ current = parent_run_id
243
+ while current is not None:
244
+ if self._top_level.get(current):
245
+ return current
246
+ # If current is a nested run, keep climbing — for simplicity return
247
+ # the immediate parent that is registered as top-level
248
+ break
249
+ return current if current in self._runs else None