Propagate agent information into LLM calls

pydantic_ai_slim/pydantic_ai/_agent_graph.py

@@ -3,14 +3,15 @@
 import asyncio
 import dataclasses
 import inspect
-import uuid
+import time
 from asyncio import Task
 from collections import defaultdict, deque
 from collections.abc import AsyncIterator, Awaitable, Callable, Iterator, Sequence
 from contextlib import asynccontextmanager, contextmanager
 from contextvars import ContextVar
 from copy import deepcopy
 from dataclasses import field, replace
+from random import Random
 from typing import TYPE_CHECKING, Any, Generic, Literal, TypeGuard, cast
 
 from opentelemetry.trace import Tracer
@@ -19,6 +20,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._ulid import ulid_as_uuid
 from pydantic_ai._utils import dataclasses_no_defaults_repr, get_union_args, now_utc
 from pydantic_ai.builtin_tools import AbstractBuiltinTool
 from pydantic_ai.capabilities.abstract import AbstractCapability
@@ -77,7 +79,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: ulid_as_uuid(Random(), lambda: time.time_ns() // 1_000_000))
     metadata: dict[str, Any] | None = None
     last_max_tokens: int | None = None
     """Last-resolved `max_tokens` from model settings, used only in error messages."""

pydantic_ai_slim/pydantic_ai/_ulid.py

@@ -0,0 +1,46 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+from random import Random
+
+
+def ulid(random: Random, ms_timestamp_generator: Callable[[], int]) -> int:
+    """Generate an integer ULID compatible with UUID v4.
+
+    ULIDs as defined by the [spec](https://github.com/ulid/spec) look like this:
+
+     01AN4Z07BY      79KA1307SR9X4MV3
+    |----------|    |----------------|
+     Timestamp         Randomness
+     48bits            80bits
+
+    In the future it would be nice to make this compatible with a UUID,
+    e.g. v4 UUIDs by setting the version and variant bits correctly.
+    We can't currently do this because setting these bits would leave us with only 7 bytes of randomness,
+    which isn't enough for the Python SDK's sampler that currently expects 8 bytes of randomness.
+    In the future OTEL will probably adopt https://www.w3.org/TR/trace-context-2/#random-trace-id-flag
+    which relies only on the lower 7 bytes of the trace ID, then all SDKs and tooling should be updated
+    and leaving only 7 bytes of randomness should be fine.
+
+    Right now we only care about:
+    - Our SDK / Python SDK's in general.
+    - The OTEL collector.
+
+    And both behave properly with 8 bytes of randomness because trace IDs were originally 64 bits
+    so to be compatible with old trace IDs nothing in OTEL can assume >8 bytes of randomness in trace IDs
+    unless they generated the trace ID themselves (e.g. the Go SDK _does_ expect >8 bytes of randomness internally).
+    """
+    # Timestamp: first 6 bytes of the ULID (48 bits)
+    # Note that it's not important that this timestamp is super precise or unique.
+    # It just needs to be roughly monotonically increasing so that the ULID is sortable, at least for our purposes.
+    timestamp = ms_timestamp_generator().to_bytes(6, byteorder='big')
+    # Randomness: next 10 bytes of the ULID (80 bits)
+    randomness = random.getrandbits(80).to_bytes(10, byteorder='big')
+    # Convert to int and return
+    return int.from_bytes(timestamp + randomness, byteorder='big')
+
+
+def ulid_as_uuid(random: Random, ms_timestamp_generator: Callable[[], int]) -> str:
+    """Generate a ULID string and return it as a hex string (without dashes) that can be used as a UUID."""
+    ulid_int = ulid(random, ms_timestamp_generator)
+    return f'{ulid_int:032x}'