diff --git a/pydantic_ai_slim/pydantic_ai/_agent_graph.py b/pydantic_ai_slim/pydantic_ai/_agent_graph.py index 37d3eb71f4..a04bdc5e87 100644 --- a/pydantic_ai_slim/pydantic_ai/_agent_graph.py +++ b/pydantic_ai_slim/pydantic_ai/_agent_graph.py @@ -3,7 +3,6 @@ from __future__ import annotations as _annotations import asyncio import dataclasses import inspect -import uuid from asyncio import Task from collections import defaultdict, deque from collections.abc import AsyncIterator, Awaitable, Callable, Iterator, Sequence @@ -20,6 +19,7 @@ from pydantic_ai._history_processor import HistoryProcessor from pydantic_ai._instrumentation import DEFAULT_INSTRUMENTATION_VERSION from pydantic_ai._tool_manager import ToolManager, ValidatedToolCall from pydantic_ai._utils import dataclasses_no_defaults_repr, get_union_args, now_utc +from pydantic_ai._uuid import uuid7 from pydantic_ai.builtin_tools import AbstractBuiltinTool from pydantic_ai.capabilities.abstract import AbstractCapability from pydantic_ai.models import ModelRequestContext @@ -77,7 +77,7 @@ class GraphAgentState: usage: _usage.RunUsage = dataclasses.field(default_factory=_usage.RunUsage) retries: int = 0 run_step: int = 0 - run_id: str = dataclasses.field(default_factory=lambda: str(uuid.uuid4())) + run_id: str = dataclasses.field(default_factory=lambda: str(uuid7())) metadata: dict[str, Any] | None = None last_max_tokens: int | None = None """Last-resolved `max_tokens` from model settings, used only in error messages.""" diff --git a/pydantic_ai_slim/pydantic_ai/_instrumentation.py b/pydantic_ai_slim/pydantic_ai/_instrumentation.py index 3c23b4475a..6e363281b7 100644 --- a/pydantic_ai_slim/pydantic_ai/_instrumentation.py +++ b/pydantic_ai_slim/pydantic_ai/_instrumentation.py @@ -1,7 +1,9 @@ from __future__ import annotations from dataclasses import dataclass -from typing import TYPE_CHECKING, ClassVar +from typing import TYPE_CHECKING, Any, ClassVar + +from opentelemetry.baggage import get_baggage if TYPE_CHECKING: from typing_extensions import Self @@ -9,6 +11,21 @@ if TYPE_CHECKING: DEFAULT_INSTRUMENTATION_VERSION = 2 """Default instrumentation version for `InstrumentationSettings`.""" +AGENT_NAME_BAGGAGE_KEY = 'gen_ai.agent.name' +RUN_ID_BAGGAGE_KEY = 'pydantic_ai.run_id' + + +def get_agent_run_baggage_attributes() -> dict[str, Any]: + """Read agent name and run ID from OTel baggage and return as span attributes.""" + attrs: dict[str, Any] = {} + agent_name = get_baggage(AGENT_NAME_BAGGAGE_KEY) + if agent_name is not None: + attrs[AGENT_NAME_BAGGAGE_KEY] = agent_name + run_id = get_baggage(RUN_ID_BAGGAGE_KEY) + if run_id is not None: + attrs[RUN_ID_BAGGAGE_KEY] = run_id + return attrs + @dataclass(frozen=True) class InstrumentationNames: diff --git a/pydantic_ai_slim/pydantic_ai/_output.py b/pydantic_ai_slim/pydantic_ai/_output.py index 87f93c4e23..4183ef6664 100644 --- a/pydantic_ai_slim/pydantic_ai/_output.py +++ b/pydantic_ai_slim/pydantic_ai/_output.py @@ -11,7 +11,7 @@ from pydantic import Json, TypeAdapter, ValidationError from pydantic_core import SchemaValidator, to_json from typing_extensions import Self, TypedDict, TypeVar -from pydantic_ai._instrumentation import InstrumentationNames +from pydantic_ai._instrumentation import InstrumentationNames, get_agent_run_baggage_attributes from pydantic_ai._utils import get_function_type_hints from . import _function_schema, _utils, messages as _messages @@ -102,6 +102,7 @@ async def execute_traced_output_function( tool_name = run_context.tool_name or getattr(function_schema.function, '__name__', 'output_function') attributes = { 'gen_ai.tool.name': tool_name, + **get_agent_run_baggage_attributes(), 'logfire.msg': f'running output function: {tool_name}', } if run_context.tool_call_id: diff --git a/pydantic_ai_slim/pydantic_ai/_tool_manager.py b/pydantic_ai_slim/pydantic_ai/_tool_manager.py index e893744102..7b268de7eb 100644 --- a/pydantic_ai_slim/pydantic_ai/_tool_manager.py +++ b/pydantic_ai_slim/pydantic_ai/_tool_manager.py @@ -13,7 +13,7 @@ from pydantic import ValidationError from typing_extensions import deprecated from . import messages as _messages -from ._instrumentation import InstrumentationNames +from ._instrumentation import InstrumentationNames, get_agent_run_baggage_attributes from ._run_context import AgentDepsT, RunContext from .exceptions import ( ApprovalRequired, @@ -528,6 +528,7 @@ class ToolManager(Generic[AgentDepsT]): # NOTE: this means `gen_ai.tool.call.id` will be included even if it was generated by pydantic-ai 'gen_ai.tool.call.id': call.tool_call_id, **({instrumentation_names.tool_arguments_attr: call.args_as_json_str()} if include_content else {}), + **get_agent_run_baggage_attributes(), 'logfire.msg': f'running tool: {call.tool_name}', # add the JSON schema so these attributes are formatted nicely in Logfire 'logfire.json_schema': json.dumps( diff --git a/pydantic_ai_slim/pydantic_ai/_uuid.py b/pydantic_ai_slim/pydantic_ai/_uuid.py new file mode 100644 index 0000000000..9e2e76a68a --- /dev/null +++ b/pydantic_ai_slim/pydantic_ai/_uuid.py @@ -0,0 +1,87 @@ +"""UUIDv7 polyfill for Python < 3.14. + +Matches the CPython 3.14 implementation: +https://github.com/python/cpython/blob/main/Lib/uuid.py + +Replace with `uuid.uuid7()` once Python 3.14 is the minimum supported version. +""" + +from __future__ import annotations + +import os +import threading +import time +import uuid + +# Global state for sub-millisecond monotonicity (Method 1, RFC 9562 §6.2). +_last_timestamp_v7: int | None = None +_last_counter_v7: int = 0 # 42-bit counter +_lock_v7 = threading.Lock() + +# version = 0b0111, variant = 0b10 +_RFC_4122_VERSION_7_FLAGS = 0x0000_0000_0000_7000_8000_0000_0000_0000 + + +def _uuid7_get_counter_and_tail() -> tuple[int, int]: + rand = int.from_bytes(os.urandom(10), 'big') + # 42-bit counter with MSB set to 0 + counter = (rand >> 32) & 0x1FF_FFFF_FFFF + # 32-bit random data + tail = rand & 0xFFFF_FFFF + return counter, tail + + +def uuid7() -> uuid.UUID: + """Generate a UUIDv7 (time-sortable UUID) per RFC 9562. + + UUIDv7 objects feature monotonicity within a millisecond. + """ + # --- 48 --- -- 4 -- --- 12 --- -- 2 -- --- 30 --- - 32 - + # unix_ts_ms | version | counter_hi | variant | counter_lo | random + # + # 'counter = counter_hi | counter_lo' is a 42-bit counter constructed + # with Method 1 of RFC 9562, §6.2, and its MSB is set to 0. + # + # 'random' is a 32-bit random value regenerated for every new UUID. + # + # If multiple UUIDs are generated within the same millisecond, the LSB + # of 'counter' is incremented by 1. When overflowing, the timestamp is + # advanced and the counter is reset to a random 42-bit integer with MSB + # set to 0. + + global _last_timestamp_v7 + global _last_counter_v7 + + nanoseconds = time.time_ns() + timestamp_ms = nanoseconds // 1_000_000 + + with _lock_v7: + if _last_timestamp_v7 is None or timestamp_ms > _last_timestamp_v7: + counter, tail = _uuid7_get_counter_and_tail() + else: + if timestamp_ms < _last_timestamp_v7: + timestamp_ms = _last_timestamp_v7 + 1 + # advance the 42-bit counter + counter = _last_counter_v7 + 1 + if counter > 0x3FF_FFFF_FFFF: + # advance the 48-bit timestamp + timestamp_ms += 1 + counter, tail = _uuid7_get_counter_and_tail() + else: + # 32-bit random data + tail = int.from_bytes(os.urandom(4), 'big') + + _last_timestamp_v7 = timestamp_ms + _last_counter_v7 = counter + + unix_ts_ms = timestamp_ms & 0xFFFF_FFFF_FFFF + counter_hi = (counter >> 30) & 0x0FFF + counter_lo = counter & 0x3FFF_FFFF + + int_uuid_7 = unix_ts_ms << 80 + int_uuid_7 |= counter_hi << 64 + int_uuid_7 |= counter_lo << 32 + int_uuid_7 |= tail & 0xFFFF_FFFF + int_uuid_7 |= _RFC_4122_VERSION_7_FLAGS + + return uuid.UUID(int=int_uuid_7) diff --git a/pydantic_ai_slim/pydantic_ai/agent/__init__.py b/pydantic_ai_slim/pydantic_ai/agent/__init__.py index 8855d01632..84241d68c0 100644 --- a/pydantic_ai_slim/pydantic_ai/agent/__init__.py +++ b/pydantic_ai_slim/pydantic_ai/agent/__init__.py @@ -13,6 +13,8 @@ from contextvars import ContextVar from pathlib import Path from typing import TYPE_CHECKING, Any, ClassVar, Literal, cast, overload +from opentelemetry.baggage import set_baggage as _otel_set_baggage +from opentelemetry.context import attach as _otel_attach, detach as _otel_detach from opentelemetry.trace import NoOpTracer, use_span from pydantic.json_schema import GenerateJsonSchema from typing_extensions import Self, TypeVar, deprecated @@ -98,21 +100,21 @@ if TYPE_CHECKING: from ..ui._web import ModelsParam __all__ = ( + 'AbstractAgent', 'Agent', + 'AgentModelSettings', 'AgentRun', 'AgentRunResult', - 'capture_run_messages', - 'EndStrategy', - 'CallToolsNode', - 'ModelRequestNode', - 'UserPromptNode', - 'InstrumentationSettings', - 'ParallelExecutionMode', - 'WrapperAgent', - 'AbstractAgent', - 'EventStreamHandler', - 'AgentModelSettings', 'BuiltinToolFunc', + 'CallToolsNode', + 'EndStrategy', + 'EventStreamHandler', + 'InstrumentationSettings', + 'ModelRequestNode', + 'ParallelExecutionMode', + 'UserPromptNode', + 'WrapperAgent', + 'capture_run_messages', ) @@ -1271,6 +1273,7 @@ class Agent(AbstractAgent[AgentDepsT, OutputDataT]): 'model_name': model_used.model_name if model_used else 'no-model', 'agent_name': agent_name, 'gen_ai.agent.name': agent_name, + 'pydantic_ai.run_id': state.run_id, 'logfire.msg': f'{agent_name} run', } if self._description is not None: @@ -1283,176 +1286,183 @@ class Agent(AbstractAgent[AgentDepsT, OutputDataT]): instrumentation_names.get_agent_run_span_name(agent_name), attributes=span_attributes, ) - run_metadata: dict[str, Any] | None = None try: - async with ( - _concurrency.get_concurrency_context(self._concurrency_limiter, f'agent:{agent_name}'), - graph.iter( - inputs=user_prompt_node, - state=state, - deps=graph_deps, - span=use_span(run_span) if run_span.is_recording() else None, - infer_name=False, - ) as graph_run, - ): - async with toolset: - agent_run = AgentRun(graph_run) - run_metadata = self._resolve_and_store_metadata(agent_run.ctx, metadata) + async with AsyncExitStack() as stack: + if run_span.is_recording(): + ctx = _otel_set_baggage('gen_ai.agent.name', agent_name) + ctx = _otel_set_baggage('pydantic_ai.run_id', state.run_id, context=ctx) + token = _otel_attach(ctx) + stack.callback(_otel_detach, token) + await stack.enter_async_context( + _concurrency.get_concurrency_context(self._concurrency_limiter, f'agent:{agent_name}') + ) + graph_run = await stack.enter_async_context( + graph.iter( + inputs=user_prompt_node, + state=state, + deps=graph_deps, + span=use_span(run_span) if run_span.is_recording() else None, + infer_name=False, + ) + ) + await stack.enter_async_context(toolset) + agent_run = AgentRun(graph_run) + run_metadata = self._resolve_and_store_metadata(agent_run.ctx, metadata) - # Build RunContext for run lifecycle hooks - run_ctx = _agent_graph.build_run_context(agent_run.ctx) + # Build RunContext for run lifecycle hooks + run_ctx = _agent_graph.build_run_context(agent_run.ctx) - # wrap_run cooperative hand-off protocol: - # - # 1. _do_run() calls before_run, sets _run_ready, then awaits _run_done. - # 2. wrap_run wraps _do_run via the capability middleware chain. - # 3. We await either _run_ready (handler started) or _wrap_task completion - # (short-circuit: wrap_run returned without calling handler). - # 4. We yield agent_run to the caller for iteration. - # 5. When the caller finishes (or an error occurs), we set _run_done. - # 6. _do_run resumes: returns the result (success) or re-raises the error. - # 7. If wrap_run catches the error and returns a recovery result, we use it. - # Otherwise the original error propagates. - _run_ready = asyncio.Event() - _run_done = asyncio.Event() - _run_error: BaseException | None = None - _wrap_context: list[tuple[ContextVar[Any], Any]] | None = None + # wrap_run cooperative hand-off protocol: + # + # 1. _do_run() calls before_run, sets _run_ready, then awaits _run_done. + # 2. wrap_run wraps _do_run via the capability middleware chain. + # 3. We await either _run_ready (handler started) or _wrap_task completion + # (short-circuit: wrap_run returned without calling handler). + # 4. We yield agent_run to the caller for iteration. + # 5. When the caller finishes (or an error occurs), we set _run_done. + # 6. _do_run resumes: returns the result (success) or re-raises the error. + # 7. If wrap_run catches the error and returns a recovery result, we use it. + # Otherwise the original error propagates. + _run_ready = asyncio.Event() + _run_done = asyncio.Event() + _run_error: BaseException | None = None + _wrap_context: list[tuple[ContextVar[Any], Any]] | None = None - async def _do_run() -> AgentRunResult[Any]: - nonlocal _wrap_context - await run_capability.before_run(run_ctx) - # Capture context vars set by wrap_run/before_run so - # they can be propagated to the outer task where - # agent_run.next() (and therefore node hooks) execute. - _current_ctx = contextvars.copy_context() - _wrap_context = [ - (var, _current_ctx[var]) - for var in _current_ctx - if var not in _outer_context or _outer_context[var] is not _current_ctx[var] - ] - _run_ready.set() - await _run_done.wait() - if _run_error is not None: - # Raise the original node error, not the potentially - # transformed version from context manager __aexit__ chains. - raise agent_run._node_error or _run_error # pyright: ignore[reportPrivateUsage] - r = agent_run.result - assert r is not None - return r + async def _do_run() -> AgentRunResult[Any]: + nonlocal _wrap_context + await run_capability.before_run(run_ctx) + # Capture context vars set by wrap_run/before_run so + # they can be propagated to the outer task where + # agent_run.next() (and therefore node hooks) execute. + _current_ctx = contextvars.copy_context() + _wrap_context = [ + (var, _current_ctx[var]) + for var in _current_ctx + if var not in _outer_context or _outer_context[var] is not _current_ctx[var] + ] + _run_ready.set() + await _run_done.wait() + if _run_error is not None: + # Raise the original node error, not the potentially + # transformed version from context manager __aexit__ chains. + raise agent_run._node_error or _run_error # pyright: ignore[reportPrivateUsage] + r = agent_run.result + assert r is not None + return r - _outer_context = contextvars.copy_context() - _wrap_task = asyncio.create_task(run_capability.wrap_run(run_ctx, handler=_do_run)) + _outer_context = contextvars.copy_context() + _wrap_task = asyncio.create_task(run_capability.wrap_run(run_ctx, handler=_do_run)) - # Wait for handler to start or wrap_run to complete (short-circuit) - _ready_waiter = asyncio.create_task(_run_ready.wait()) - await asyncio.wait({_ready_waiter, _wrap_task}, return_when=asyncio.FIRST_COMPLETED) - _ready_waiter.cancel() + # Wait for handler to start or wrap_run to complete (short-circuit) + _ready_waiter = asyncio.create_task(_run_ready.wait()) + await asyncio.wait({_ready_waiter, _wrap_task}, return_when=asyncio.FIRST_COMPLETED) + _ready_waiter.cancel() - # Propagate context vars set by wrap_run/before_run to - # the outer task so that agent_run.next() (and therefore - # node hooks) can see them. - _context_tokens: list[tuple[ContextVar[Any], contextvars.Token[Any]]] = [] - # Note: indexing instead of tuple unpacking because pyright - # can't resolve types through nonlocal + Optional unpacking. - for _cv_pair in _wrap_context or (): - _context_tokens.append((_cv_pair[0], _cv_pair[0].set(_cv_pair[1]))) + # Propagate context vars set by wrap_run/before_run to + # the outer task so that agent_run.next() (and therefore + # node hooks) can see them. + _context_tokens: list[tuple[ContextVar[Any], contextvars.Token[Any]]] = [] + # Note: indexing instead of tuple unpacking because pyright + # can't resolve types through nonlocal + Optional unpacking. + for _cv_pair in _wrap_context or (): + _context_tokens.append((_cv_pair[0], _cv_pair[0].set(_cv_pair[1]))) - async def _finalize_result(r: AgentRunResult[Any]) -> None: - """Call after_run, store the result override, and clear any pending error.""" - nonlocal _run_error - r = await run_capability.after_run(run_ctx, result=r) - agent_run._result_override = r # pyright: ignore[reportPrivateUsage] - _run_error = None + async def _finalize_result(r: AgentRunResult[Any]) -> None: + """Call after_run, store the result override, and clear any pending error.""" + nonlocal _run_error + r = await run_capability.after_run(run_ctx, result=r) + agent_run._result_override = r # pyright: ignore[reportPrivateUsage] + _run_error = None + + try: + _short_circuited = _wrap_task.done() and not _run_ready.is_set() + if _short_circuited: + await _finalize_result(_wrap_task.result()) try: - _short_circuited = _wrap_task.done() and not _run_ready.is_set() - if _short_circuited: - await _finalize_result(_wrap_task.result()) - - try: - yield agent_run - except BaseException as _exc: - # Use the original node error if available, since context manager - # __aexit__ chains (GraphRun → anyio TaskGroup) may transform - # the exception (e.g. into CancelledError or ExceptionGroup). - _run_error = agent_run._node_error or _exc # pyright: ignore[reportPrivateUsage] - # Don't attempt recovery for GeneratorExit/KeyboardInterrupt — - # awaiting _wrap_task during cleanup could delay shutdown. - if isinstance(_run_error, (GeneratorExit, KeyboardInterrupt)): - raise - # Don't re-raise yet — give wrap_run a chance to recover. - # If wrap_run catches the error from handler() and returns - # a recovery result, the exception will be suppressed. - finally: - if agent_run.result is not None: - run_metadata = self._resolve_and_store_metadata(agent_run.ctx, metadata) - else: - run_metadata = graph_run.state.metadata - - if not _short_circuited: - _run_done.set() - if _run_error is None and agent_run.result is not None: - await _finalize_result(await _wrap_task) - elif _run_error is not None: - # Error path: await wrap_run to see if it recovers. - # _do_run() re-raises _run_error; if wrap_run catches - # it and returns a result, recovery succeeds. - try: - await _finalize_result(await _wrap_task) - except BaseException as _wrap_exc: - # Attach wrap_run's own errors as context so they're - # visible in tracebacks (but don't mask the original). - # Skip CancelledError: it's expected cancellation propagation, - # and setting __context__ on it causes hangs on Python 3.10. - if ( - not isinstance(_wrap_exc, asyncio.CancelledError) - and _wrap_exc is not _run_error - ): - _run_error.__context__ = _wrap_exc # pragma: no cover — only fires for bugs in wrap_run implementations - elif ( - not _wrap_task.done() - ): # pragma: no branch — _run_done.set() can't complete _wrap_task synchronously - _wrap_task.cancel() - try: - await _wrap_task - except (asyncio.CancelledError, BaseException): - pass - - # If wrap_run didn't recover, give on_run_error a chance. - if _run_error is not None: - try: - _result = await run_capability.on_run_error(run_ctx, error=_run_error) - except BaseException as _on_error_exc: - _run_error = _on_error_exc - else: - await _finalize_result(_result) - - # If on_run_error didn't recover either, re-raise. - # In an @asynccontextmanager, not re-raising suppresses the exception. - if _run_error is not None: - raise _run_error + yield agent_run + except BaseException as _exc: + # Use the original node error if available, since context manager + # __aexit__ chains (GraphRun → anyio TaskGroup) may transform + # the exception (e.g. into CancelledError or ExceptionGroup). + _run_error = agent_run._node_error or _exc # pyright: ignore[reportPrivateUsage] + # Don't attempt recovery for GeneratorExit/KeyboardInterrupt — + # awaiting _wrap_task during cleanup could delay shutdown. + if isinstance(_run_error, (GeneratorExit, KeyboardInterrupt)): + raise + # Don't re-raise yet — give wrap_run a chance to recover. + # If wrap_run catches the error from handler() and returns + # a recovery result, the exception will be suppressed. finally: - # Always restore context vars, even on - # GeneratorExit/KeyboardInterrupt. - for _var, _token in _context_tokens: - _var.reset(_token) + if agent_run.result is not None: + run_metadata = self._resolve_and_store_metadata(agent_run.ctx, metadata) + else: + run_metadata = graph_run.state.metadata - final_result = agent_run.result - if ( - instrumentation_settings - and instrumentation_settings.include_content - and run_span.is_recording() - and final_result is not None - ): - run_span.set_attribute( - 'final_result', - ( - final_result.output - if isinstance(final_result.output, str) - else json.dumps(InstrumentedModel.serialize_any(final_result.output)) - ), - ) + if not _short_circuited: + _run_done.set() + if _run_error is None and agent_run.result is not None: + await _finalize_result(await _wrap_task) + elif _run_error is not None: + # Error path: await wrap_run to see if it recovers. + # _do_run() re-raises _run_error; if wrap_run catches + # it and returns a result, recovery succeeds. + try: + await _finalize_result(await _wrap_task) + except BaseException as _wrap_exc: + # Attach wrap_run's own errors as context so they're + # visible in tracebacks (but don't mask the original). + # Skip CancelledError: it's expected cancellation propagation, + # and setting __context__ on it causes hangs on Python 3.10. + if ( + not isinstance(_wrap_exc, asyncio.CancelledError) + and _wrap_exc is not _run_error + ): + _run_error.__context__ = _wrap_exc # pragma: no cover — only fires for bugs in wrap_run implementations + elif ( + not _wrap_task.done() + ): # pragma: no branch — _run_done.set() can't complete _wrap_task synchronously + _wrap_task.cancel() + try: + await _wrap_task + except (asyncio.CancelledError, BaseException): + pass + + # If wrap_run didn't recover, give on_run_error a chance. + if _run_error is not None: + try: + _result = await run_capability.on_run_error(run_ctx, error=_run_error) + except BaseException as _on_error_exc: + _run_error = _on_error_exc + else: + await _finalize_result(_result) + + # If on_run_error didn't recover either, re-raise. + # In an @asynccontextmanager, not re-raising suppresses the exception. + if _run_error is not None: + raise _run_error + finally: + # Always restore context vars, even on + # GeneratorExit/KeyboardInterrupt. + for _var, _token in _context_tokens: + _var.reset(_token) + + final_result = agent_run.result + if ( + instrumentation_settings + and instrumentation_settings.include_content + and run_span.is_recording() + and final_result is not None + ): + run_span.set_attribute( + 'final_result', + ( + final_result.output + if isinstance(final_result.output, str) + else json.dumps(InstrumentedModel.serialize_any(final_result.output)) + ), + ) finally: try: if instrumentation_settings and run_span.is_recording(): diff --git a/pydantic_ai_slim/pydantic_ai/models/instrumented.py b/pydantic_ai_slim/pydantic_ai/models/instrumented.py index 902b812b95..a765354402 100644 --- a/pydantic_ai_slim/pydantic_ai/models/instrumented.py +++ b/pydantic_ai_slim/pydantic_ai/models/instrumented.py @@ -21,7 +21,7 @@ from opentelemetry.trace import Span, SpanKind, Tracer, TracerProvider, get_trac from opentelemetry.util.types import AttributeValue from pydantic import TypeAdapter -from pydantic_ai._instrumentation import DEFAULT_INSTRUMENTATION_VERSION +from pydantic_ai._instrumentation import DEFAULT_INSTRUMENTATION_VERSION, get_agent_run_baggage_attributes from .. import _otel_messages from .._run_context import RunContext @@ -456,6 +456,7 @@ class InstrumentedModel(WrapperModel): 'gen_ai.operation.name': operation, **self.model_attributes(self.wrapped), **self.model_request_parameters_attributes(model_request_parameters), + **get_agent_run_baggage_attributes(), 'logfire.json_schema': json.dumps( { 'type': 'object', diff --git a/tests/models/test_fallback.py b/tests/models/test_fallback.py index a0c8c807cb..6bb68fb67a 100644 --- a/tests/models/test_fallback.py +++ b/tests/models/test_fallback.py @@ -172,6 +172,8 @@ def test_first_failed_instrumented(capfire: CaptureLogfire) -> None: 'thinking': None, }, 'logfire.span_type': 'span', + 'gen_ai.agent.name': 'agent', + 'pydantic_ai.run_id': IsStr(), 'gen_ai.provider.name': 'function', 'logfire.msg': 'chat fallback:function:failure_response:,function:success_response:', 'gen_ai.system': 'function', @@ -203,6 +205,7 @@ def test_first_failed_instrumented(capfire: CaptureLogfire) -> None: 'model_name': 'fallback:function:failure_response:,function:success_response:', 'agent_name': 'agent', 'gen_ai.agent.name': 'agent', + 'pydantic_ai.run_id': IsStr(), 'logfire.msg': 'agent run', 'logfire.span_type': 'span', 'gen_ai.usage.input_tokens': 51, @@ -283,6 +286,8 @@ async def test_first_failed_instrumented_stream(capfire: CaptureLogfire) -> None 'thinking': None, }, 'logfire.span_type': 'span', + 'gen_ai.agent.name': 'agent', + 'pydantic_ai.run_id': IsStr(), 'gen_ai.provider.name': 'function', 'logfire.msg': 'chat fallback:function::failure_response_stream,function::success_response_stream', 'gen_ai.system': 'function', @@ -314,6 +319,7 @@ async def test_first_failed_instrumented_stream(capfire: CaptureLogfire) -> None 'model_name': 'fallback:function::failure_response_stream,function::success_response_stream', 'agent_name': 'agent', 'gen_ai.agent.name': 'agent', + 'pydantic_ai.run_id': IsStr(), 'logfire.msg': 'agent run', 'logfire.span_type': 'span', 'final_result': 'hello world', @@ -401,6 +407,8 @@ def test_all_failed_instrumented(capfire: CaptureLogfire) -> None: }, 'logfire.span_type': 'span', 'logfire.msg': 'chat fallback:function:failure_response:,function:failure_response:', + 'gen_ai.agent.name': 'agent', + 'pydantic_ai.run_id': IsStr(), 'logfire.level_num': 17, 'gen_ai.response.model': 'fallback:function:failure_response:,function:failure_response:', }, @@ -427,6 +435,7 @@ def test_all_failed_instrumented(capfire: CaptureLogfire) -> None: 'model_name': 'fallback:function:failure_response:,function:failure_response:', 'agent_name': 'agent', 'gen_ai.agent.name': 'agent', + 'pydantic_ai.run_id': IsStr(), 'logfire.msg': 'agent run', 'logfire.span_type': 'span', 'logfire.exception.fingerprint': '0000000000000000000000000000000000000000000000000000000000000000', @@ -894,6 +903,8 @@ Don't include any text or Markdown fencing before or after. 'thinking': None, }, 'logfire.span_type': 'span', + 'gen_ai.agent.name': 'agent', + 'pydantic_ai.run_id': IsStr(), 'gen_ai.provider.name': 'function', 'logfire.msg': 'chat fallback:function:tool_output_func:,function:prompted_output_func:', 'gen_ai.system': 'function', @@ -927,6 +938,7 @@ Don't include any text or Markdown fencing before or after. 'model_name': 'fallback:function:tool_output_func:,function:prompted_output_func:', 'agent_name': 'agent', 'gen_ai.agent.name': 'agent', + 'pydantic_ai.run_id': IsStr(), 'logfire.msg': 'agent run', 'logfire.span_type': 'span', 'gen_ai.usage.input_tokens': 51, diff --git a/tests/test_logfire.py b/tests/test_logfire.py index da607b7e0e..7a4f2737a8 100644 --- a/tests/test_logfire.py +++ b/tests/test_logfire.py @@ -237,6 +237,7 @@ def test_logfire( 'model_name': 'test', 'agent_name': 'my_agent', 'gen_ai.agent.name': 'my_agent', + 'pydantic_ai.run_id': IsStr(), 'logfire.msg': 'my_agent run', 'logfire.span_type': 'span', 'final_result': '{"my_ret":"1"}', @@ -293,6 +294,7 @@ def test_logfire( 'model_name': 'test', 'agent_name': 'my_agent', 'gen_ai.agent.name': 'my_agent', + 'pydantic_ai.run_id': IsStr(), 'logfire.msg': 'my_agent run', 'logfire.span_type': 'span', 'gen_ai.usage.input_tokens': 103, @@ -530,6 +532,8 @@ def test_logfire( 'gen_ai.tool.definitions': '[{"type": "function", "name": "my_ret", "parameters": {"additionalProperties": false, "properties": {"x": {"type": "integer"}}, "required": ["x"], "type": "object"}}]', 'logfire.json_schema': IsJson(), 'logfire.span_type': 'span', + 'gen_ai.agent.name': 'my_agent', + 'pydantic_ai.run_id': IsStr(), 'logfire.msg': 'chat test', 'gen_ai.response.model': 'test', 'gen_ai.usage.input_tokens': 51, @@ -597,6 +601,7 @@ def test_instructions_with_structured_output( 'model_name': 'test', 'agent_name': 'my_agent', 'gen_ai.agent.name': 'my_agent', + 'pydantic_ai.run_id': IsStr(), 'logfire.msg': 'my_agent run', 'logfire.span_type': 'span', 'gen_ai.usage.input_tokens': 51, @@ -716,6 +721,7 @@ def test_instructions_with_structured_output( 'model_name': 'test', 'agent_name': 'my_agent', 'gen_ai.agent.name': 'my_agent', + 'pydantic_ai.run_id': IsStr(), 'logfire.msg': 'my_agent run', 'logfire.span_type': 'span', 'final_result': '{"content": "a"}', @@ -807,6 +813,7 @@ def test_instructions_with_structured_output_exclude_content(get_logfire_summary 'model_name': 'test', 'agent_name': 'my_agent', 'gen_ai.agent.name': 'my_agent', + 'pydantic_ai.run_id': IsStr(), 'logfire.msg': 'my_agent run', 'logfire.span_type': 'span', 'gen_ai.usage.input_tokens': 51, @@ -925,6 +932,7 @@ def test_instructions_with_structured_output_exclude_content_v2_v3( 'model_name': 'test', 'agent_name': 'my_agent', 'gen_ai.agent.name': 'my_agent', + 'pydantic_ai.run_id': IsStr(), 'logfire.msg': 'my_agent run', 'logfire.span_type': 'span', 'gen_ai.usage.input_tokens': 51, @@ -1012,6 +1020,8 @@ def test_instructions_with_structured_output_exclude_content_v2_v3( 'gen_ai.tool.definitions': '[{"type": "function", "name": "final_result", "description": "The final response which ends this conversation", "parameters": {"properties": {"content": {"type": "string"}}, "required": ["content"], "title": "MyOutput", "type": "object"}}]', 'logfire.span_type': 'span', 'logfire.msg': 'chat test', + 'gen_ai.agent.name': 'my_agent', + 'pydantic_ai.run_id': IsStr(), 'gen_ai.input.messages': IsJson(snapshot([{'role': 'user', 'parts': [{'type': 'text'}]}])), 'gen_ai.output.messages': IsJson( snapshot( @@ -1102,6 +1112,7 @@ async def test_aggregated_usage_attribute_names(capfire: CaptureLogfire) -> None 'model_name': 'function:model_function:', 'agent_name': 'agent', 'gen_ai.agent.name': 'agent', + 'pydantic_ai.run_id': IsStr(), 'logfire.msg': 'agent run', 'logfire.span_type': 'span', 'final_result': 'Hello!', @@ -1168,6 +1179,8 @@ async def test_feedback(capfire: CaptureLogfire) -> None: }, 'logfire.span_type': 'span', 'logfire.msg': 'chat test', + 'gen_ai.agent.name': 'agent', + 'pydantic_ai.run_id': IsStr(), 'gen_ai.usage.input_tokens': 51, 'gen_ai.usage.output_tokens': 4, 'gen_ai.response.model': 'test', @@ -1208,6 +1221,7 @@ async def test_feedback(capfire: CaptureLogfire) -> None: 'model_name': 'test', 'agent_name': 'agent', 'gen_ai.agent.name': 'agent', + 'pydantic_ai.run_id': IsStr(), 'logfire.msg': 'agent run', 'logfire.span_type': 'span', 'gen_ai.usage.input_tokens': 51, @@ -1315,6 +1329,8 @@ def test_include_tool_args_span_attributes( ) ), 'logfire.span_type': 'span', + 'gen_ai.agent.name': 'my_agent', + 'pydantic_ai.run_id': IsStr(), 'tool_response': """\ Tool error @@ -1344,6 +1360,8 @@ Fix the errors and try again.\ } ) ), + 'gen_ai.agent.name': 'my_agent', + 'pydantic_ai.run_id': IsStr(), 'logfire.span_type': 'span', } ) @@ -1366,6 +1384,8 @@ Fix the errors and try again.\ ) ), 'logfire.span_type': 'span', + 'gen_ai.agent.name': 'my_agent', + 'pydantic_ai.run_id': IsStr(), 'logfire.level_num': 17, } ) @@ -1387,6 +1407,8 @@ Fix the errors and try again.\ ) ), 'logfire.span_type': 'span', + 'gen_ai.agent.name': 'my_agent', + 'pydantic_ai.run_id': IsStr(), } ) @@ -1468,6 +1490,8 @@ def test_logfire_output_function_v2_v3( ) ), 'logfire.span_type': 'span', + 'gen_ai.agent.name': 'my_agent', + 'pydantic_ai.run_id': IsStr(), 'tool_response': '{"temperature": 28.7, "description": "sunny"}', } ) @@ -1515,6 +1539,8 @@ def test_logfire_output_function_v2_v3( ) ), 'logfire.span_type': 'span', + 'gen_ai.agent.name': 'my_agent', + 'pydantic_ai.run_id': IsStr(), 'gen_ai.tool.call.result': '{"temperature": 28.7, "description": "sunny"}', } ) @@ -1568,6 +1594,8 @@ def test_output_type_function_logfire_attributes( ) ), 'logfire.span_type': 'span', + 'gen_ai.agent.name': 'my_agent', + 'pydantic_ai.run_id': IsStr(), 'tool_response': '{"temperature": 28.7, "description": "sunny"}', } ) @@ -1579,6 +1607,8 @@ def test_output_type_function_logfire_attributes( 'logfire.msg': 'running output function: final_result', 'logfire.json_schema': '{"type": "object", "properties": {"gen_ai.tool.name": {}, "gen_ai.tool.call.id": {}}}', 'logfire.span_type': 'span', + 'gen_ai.agent.name': 'my_agent', + 'pydantic_ai.run_id': IsStr(), } ) @@ -1632,6 +1662,8 @@ def test_output_type_function_with_run_context_logfire_attributes( ) ), 'logfire.span_type': 'span', + 'gen_ai.agent.name': 'my_agent', + 'pydantic_ai.run_id': IsStr(), 'tool_response': '{"temperature": 28.7, "description": "sunny"}', } ) @@ -1643,6 +1675,8 @@ def test_output_type_function_with_run_context_logfire_attributes( 'logfire.msg': 'running output function: final_result', 'logfire.json_schema': '{"type": "object", "properties": {"gen_ai.tool.name": {}, "gen_ai.tool.call.id": {}}}', 'logfire.span_type': 'span', + 'gen_ai.agent.name': 'my_agent', + 'pydantic_ai.run_id': IsStr(), } ) @@ -1704,6 +1738,8 @@ def test_output_type_function_with_retry_logfire_attributes( ) ), 'logfire.span_type': 'span', + 'gen_ai.agent.name': 'my_agent', + 'pydantic_ai.run_id': IsStr(), 'logfire.level_num': 17, }, { @@ -1725,6 +1761,8 @@ def test_output_type_function_with_retry_logfire_attributes( ) ), 'logfire.span_type': 'span', + 'gen_ai.agent.name': 'my_agent', + 'pydantic_ai.run_id': IsStr(), 'tool_response': '{"temperature": 28.7, "description": "sunny"}', }, ] @@ -1740,6 +1778,8 @@ def test_output_type_function_with_retry_logfire_attributes( snapshot({'type': 'object', 'properties': {'gen_ai.tool.name': {}, 'gen_ai.tool.call.id': {}}}) ), 'logfire.span_type': 'span', + 'gen_ai.agent.name': 'my_agent', + 'pydantic_ai.run_id': IsStr(), 'logfire.level_num': 17, }, { @@ -1750,6 +1790,8 @@ def test_output_type_function_with_retry_logfire_attributes( snapshot({'type': 'object', 'properties': {'gen_ai.tool.name': {}, 'gen_ai.tool.call.id': {}}}) ), 'logfire.span_type': 'span', + 'gen_ai.agent.name': 'my_agent', + 'pydantic_ai.run_id': IsStr(), }, ] ) @@ -1802,6 +1844,8 @@ def test_output_type_function_with_custom_tool_name_logfire_attributes( ) ), 'logfire.span_type': 'span', + 'gen_ai.agent.name': 'my_agent', + 'pydantic_ai.run_id': IsStr(), 'tool_response': '{"temperature": 28.7, "description": "sunny"}', } ) @@ -1815,6 +1859,8 @@ def test_output_type_function_with_custom_tool_name_logfire_attributes( snapshot({'type': 'object', 'properties': {'gen_ai.tool.name': {}, 'gen_ai.tool.call.id': {}}}) ), 'logfire.span_type': 'span', + 'gen_ai.agent.name': 'my_agent', + 'pydantic_ai.run_id': IsStr(), } ) @@ -1873,6 +1919,8 @@ def test_output_type_bound_instance_method_logfire_attributes( ) ), 'logfire.span_type': 'span', + 'gen_ai.agent.name': 'my_agent', + 'pydantic_ai.run_id': IsStr(), 'tool_response': '{"temperature": 28.7, "description": "sunny"}', } ) @@ -1886,6 +1934,8 @@ def test_output_type_bound_instance_method_logfire_attributes( snapshot({'type': 'object', 'properties': {'gen_ai.tool.name': {}, 'gen_ai.tool.call.id': {}}}) ), 'logfire.span_type': 'span', + 'gen_ai.agent.name': 'my_agent', + 'pydantic_ai.run_id': IsStr(), } ) @@ -1945,6 +1995,8 @@ def test_output_type_bound_instance_method_with_run_context_logfire_attributes( ) ), 'logfire.span_type': 'span', + 'gen_ai.agent.name': 'my_agent', + 'pydantic_ai.run_id': IsStr(), 'tool_response': '{"temperature": 28.7, "description": "sunny"}', } ) @@ -1958,6 +2010,8 @@ def test_output_type_bound_instance_method_with_run_context_logfire_attributes( snapshot({'type': 'object', 'properties': {'gen_ai.tool.name': {}, 'gen_ai.tool.call.id': {}}}) ), 'logfire.span_type': 'span', + 'gen_ai.agent.name': 'my_agent', + 'pydantic_ai.run_id': IsStr(), } ) @@ -2012,6 +2066,8 @@ def test_output_type_async_function_logfire_attributes( ) ), 'logfire.span_type': 'span', + 'gen_ai.agent.name': 'my_agent', + 'pydantic_ai.run_id': IsStr(), 'tool_response': '{"temperature": 28.7, "description": "sunny"}', } ) @@ -2025,6 +2081,8 @@ def test_output_type_async_function_logfire_attributes( snapshot({'type': 'object', 'properties': {'gen_ai.tool.name': {}, 'gen_ai.tool.call.id': {}}}) ), 'logfire.span_type': 'span', + 'gen_ai.agent.name': 'my_agent', + 'pydantic_ai.run_id': IsStr(), } ) @@ -2082,6 +2140,8 @@ def test_text_output_function_logfire_attributes( ) ), 'logfire.span_type': 'span', + 'gen_ai.agent.name': 'my_agent', + 'pydantic_ai.run_id': IsStr(), 'tool_response': 'HELLO WORLD', } ) @@ -2092,6 +2152,8 @@ def test_text_output_function_logfire_attributes( 'logfire.msg': 'running output function: upcase_text', 'logfire.json_schema': IsJson(snapshot({'type': 'object', 'properties': {'gen_ai.tool.name': {}}})), 'logfire.span_type': 'span', + 'gen_ai.agent.name': 'my_agent', + 'pydantic_ai.run_id': IsStr(), } ) @@ -2154,6 +2216,8 @@ def test_prompted_output_function_logfire_attributes( ) ), 'logfire.span_type': 'span', + 'gen_ai.agent.name': 'agent', + 'pydantic_ai.run_id': IsStr(), 'tool_response': 'HELLO WORLD', } ) @@ -2164,6 +2228,8 @@ def test_prompted_output_function_logfire_attributes( 'logfire.msg': 'running output function: upcase_text', 'logfire.json_schema': IsJson(snapshot({'type': 'object', 'properties': {'gen_ai.tool.name': {}}})), 'logfire.span_type': 'span', + 'gen_ai.agent.name': 'agent', + 'pydantic_ai.run_id': IsStr(), } ) @@ -2226,6 +2292,8 @@ def test_output_type_text_output_function_with_retry_logfire_attributes( ) ), 'logfire.span_type': 'span', + 'gen_ai.agent.name': 'my_agent', + 'pydantic_ai.run_id': IsStr(), 'logfire.level_num': 17, }, { @@ -2245,6 +2313,8 @@ def test_output_type_text_output_function_with_retry_logfire_attributes( ) ), 'logfire.span_type': 'span', + 'gen_ai.agent.name': 'my_agent', + 'pydantic_ai.run_id': IsStr(), 'tool_response': '{"temperature": 28.7, "description": "sunny"}', }, ] @@ -2257,6 +2327,8 @@ def test_output_type_text_output_function_with_retry_logfire_attributes( 'logfire.msg': 'running output function: get_weather_with_retry', 'logfire.json_schema': IsJson(snapshot({'type': 'object', 'properties': {'gen_ai.tool.name': {}}})), 'logfire.span_type': 'span', + 'gen_ai.agent.name': 'my_agent', + 'pydantic_ai.run_id': IsStr(), 'logfire.level_num': 17, }, { @@ -2264,6 +2336,8 @@ def test_output_type_text_output_function_with_retry_logfire_attributes( 'logfire.msg': 'running output function: get_weather_with_retry', 'logfire.json_schema': IsJson(snapshot({'type': 'object', 'properties': {'gen_ai.tool.name': {}}})), 'logfire.span_type': 'span', + 'gen_ai.agent.name': 'my_agent', + 'pydantic_ai.run_id': IsStr(), }, ] ) @@ -2298,6 +2372,7 @@ def test_static_function_instructions_in_agent_run_span( 'model_name': 'test', 'agent_name': 'my_agent', 'gen_ai.agent.name': 'my_agent', + 'pydantic_ai.run_id': IsStr(), 'logfire.msg': 'my_agent run', 'logfire.span_type': 'span', 'gen_ai.usage.input_tokens': 51, @@ -2417,6 +2492,7 @@ def test_static_function_instructions_in_agent_run_span( 'model_name': 'test', 'agent_name': 'my_agent', 'gen_ai.agent.name': 'my_agent', + 'pydantic_ai.run_id': IsStr(), 'logfire.msg': 'my_agent run', 'logfire.span_type': 'span', 'final_result': '{"content": "a"}', @@ -2522,6 +2598,7 @@ def test_dynamic_function_instructions_in_agent_run_span( 'model_name': 'test', 'agent_name': 'my_agent', 'gen_ai.agent.name': 'my_agent', + 'pydantic_ai.run_id': IsStr(), 'logfire.msg': 'my_agent run', 'logfire.span_type': 'span', 'gen_ai.usage.input_tokens': 107, @@ -2669,6 +2746,7 @@ def test_dynamic_function_instructions_in_agent_run_span( 'model_name': 'test', 'agent_name': 'my_agent', 'gen_ai.agent.name': 'my_agent', + 'pydantic_ai.run_id': IsStr(), 'logfire.msg': 'my_agent run', 'logfire.span_type': 'span', 'final_result': '{"content": "a"}', @@ -2805,6 +2883,7 @@ def test_function_instructions_with_history_in_agent_run_span( 'model_name': 'test', 'agent_name': 'my_agent', 'gen_ai.agent.name': 'my_agent', + 'pydantic_ai.run_id': IsStr(), 'logfire.msg': 'my_agent run', 'logfire.span_type': 'span', 'gen_ai.usage.input_tokens': 52, @@ -2950,6 +3029,7 @@ def test_function_instructions_with_history_in_agent_run_span( 'model_name': 'test', 'agent_name': 'my_agent', 'gen_ai.agent.name': 'my_agent', + 'pydantic_ai.run_id': IsStr(), 'logfire.msg': 'my_agent run', 'logfire.span_type': 'span', 'final_result': '{"content": "a"}', @@ -3058,6 +3138,7 @@ async def test_run_stream( 'model_name': 'test', 'agent_name': 'my_agent', 'gen_ai.agent.name': 'my_agent', + 'pydantic_ai.run_id': IsStr(), 'logfire.msg': 'my_agent run', 'logfire.span_type': 'span', 'final_result': 'success (no tool calls)', @@ -3154,6 +3235,7 @@ async def test_run_stream( 'model_name': 'test', 'agent_name': 'my_agent', 'gen_ai.agent.name': 'my_agent', + 'pydantic_ai.run_id': IsStr(), 'logfire.msg': 'my_agent run', 'logfire.span_type': 'span', 'final_result': 'success (no tool calls)', @@ -3244,6 +3326,8 @@ def test_deferral_call_deferred_v2(capfire: CaptureLogfire) -> None: }, }, 'logfire.span_type': 'span', + 'gen_ai.agent.name': 'agent', + 'pydantic_ai.run_id': IsStr(), 'pydantic_ai.tool.deferral.name': 'CallDeferred', 'pydantic_ai.tool.deferral.metadata': {'task_id': 'task-123'}, 'logfire.level_num': 17, @@ -3301,6 +3385,8 @@ def test_deferral_approval_required_v2(capfire: CaptureLogfire) -> None: }, }, 'logfire.span_type': 'span', + 'gen_ai.agent.name': 'agent', + 'pydantic_ai.run_id': IsStr(), 'pydantic_ai.tool.deferral.name': 'ApprovalRequired', 'pydantic_ai.tool.deferral.metadata': {'task_id': 'task-123'}, 'logfire.level_num': 17, @@ -3358,6 +3444,8 @@ def test_deferral_call_deferred_v5(capfire: CaptureLogfire) -> None: }, }, 'logfire.span_type': 'span', + 'gen_ai.agent.name': 'agent', + 'pydantic_ai.run_id': IsStr(), 'pydantic_ai.tool.deferral.name': 'CallDeferred', 'pydantic_ai.tool.deferral.metadata': {'task_id': 'task-123'}, }, @@ -3402,6 +3490,8 @@ def test_deferral_approval_required_v5(capfire: CaptureLogfire) -> None: }, }, 'logfire.span_type': 'span', + 'gen_ai.agent.name': 'agent', + 'pydantic_ai.run_id': IsStr(), 'pydantic_ai.tool.deferral.name': 'ApprovalRequired', 'pydantic_ai.tool.deferral.metadata': {'task_id': 'task-123'}, }, @@ -3448,6 +3538,8 @@ def test_deferral_no_metadata(capfire: CaptureLogfire) -> None: }, }, 'logfire.span_type': 'span', + 'gen_ai.agent.name': 'agent', + 'pydantic_ai.run_id': IsStr(), 'pydantic_ai.tool.deferral.name': 'CallDeferred', }, } @@ -3498,6 +3590,8 @@ def test_deferral_non_serializable_metadata(capfire: CaptureLogfire) -> None: }, }, 'logfire.span_type': 'span', + 'gen_ai.agent.name': 'agent', + 'pydantic_ai.run_id': IsStr(), 'pydantic_ai.tool.deferral.name': 'CallDeferred', 'pydantic_ai.tool.deferral.metadata': "{'obj': }", }, diff --git a/tests/test_uuid.py b/tests/test_uuid.py new file mode 100644 index 0000000000..4571c29479 --- /dev/null +++ b/tests/test_uuid.py @@ -0,0 +1,61 @@ +"""Tests for the UUIDv7 polyfill.""" + +from unittest.mock import patch + +import pydantic_ai._uuid as uuid_mod +from pydantic_ai._uuid import uuid7 + + +def test_uuid7_basic() -> None: + u = uuid7() + assert u.version == 7 + assert u.variant == 'specified in RFC 4122' + assert len(str(u)) == 36 + + +def test_uuid7_monotonic_within_millisecond() -> None: + """Multiple uuid7() calls in the same millisecond produce monotonically increasing values.""" + ids = [uuid7() for _ in range(100)] + for a, b in zip(ids, ids[1:]): + assert a < b + + +def test_uuid7_clock_regression() -> None: + """uuid7() handles system clock going backward (e.g. NTP adjustment).""" + saved_ts = uuid_mod._last_timestamp_v7 # pyright: ignore[reportPrivateUsage] + saved_counter = uuid_mod._last_counter_v7 # pyright: ignore[reportPrivateUsage] + try: + # Generate a UUID at a known timestamp + with patch.object(uuid_mod.time, 'time_ns', return_value=2_000 * 1_000_000_000): + u1 = uuid7() + + # Clock goes backward + with patch.object(uuid_mod.time, 'time_ns', return_value=1_000 * 1_000_000_000): + u2 = uuid7() + + assert u2 > u1 + assert u2.version == 7 + finally: + uuid_mod._last_timestamp_v7 = saved_ts # pyright: ignore[reportPrivateUsage] + uuid_mod._last_counter_v7 = saved_counter # pyright: ignore[reportPrivateUsage] + + +def test_uuid7_counter_overflow() -> None: + """uuid7() advances the timestamp when the 42-bit counter overflows.""" + saved_ts = uuid_mod._last_timestamp_v7 # pyright: ignore[reportPrivateUsage] + saved_counter = uuid_mod._last_counter_v7 # pyright: ignore[reportPrivateUsage] + try: + # Set counter to max so the next same-ms call overflows + with patch.object(uuid_mod.time, 'time_ns', return_value=3_000 * 1_000_000_000): + u1 = uuid7() + + uuid_mod._last_counter_v7 = 0x3FF_FFFF_FFFF # pyright: ignore[reportPrivateUsage] + + with patch.object(uuid_mod.time, 'time_ns', return_value=3_000 * 1_000_000_000): + u2 = uuid7() + + assert u2 > u1 + assert u2.version == 7 + finally: + uuid_mod._last_timestamp_v7 = saved_ts # pyright: ignore[reportPrivateUsage] + uuid_mod._last_counter_v7 = saved_counter # pyright: ignore[reportPrivateUsage]