feat(observability): Enhance prompt context injection and identity tr…

.github/workflows/python-tests.yml

@@ -15,7 +15,7 @@ jobs:
 
     strategy:
       matrix:
-        python-version: ['3.8', '3.9', '3.10', '3.11', '3.12', '3.13']
+        python-version: ['3.10', '3.11', '3.12', '3.13', '3.14']
 
     steps:
     - uses: actions/checkout@v3
@@ -29,9 +29,8 @@ jobs:
       working-directory: ./
       run: |
         python -m pip install --upgrade pip
-        pip install pytest pytest-cov
+        pip install pytest pytest-cov anyio pytest-asyncio
         pip install -e .
-        pip install -r dev-requirements.txt
 
     - name: Run tests
       working-directory: ./

basalt/observability/api.py

@@ -4,7 +4,7 @@
 import inspect
 from collections.abc import Callable, Sequence
 from contextlib import ContextDecorator
-from typing import TYPE_CHECKING, Any, NotRequired, TypedDict, TypeVar
+from typing import TYPE_CHECKING, Any, TypeVar
 
 from opentelemetry.trace import StatusCode
 
@@ -31,6 +31,7 @@
     get_root_span_handle,
 )
 from .decorators import ObserveKind
+from .types import Identity
 from .utils import (
     apply_llm_request_metadata,
     apply_llm_response_metadata,
@@ -49,28 +50,6 @@
 F = TypeVar("F", bound=Callable[..., Any])
 
 
-class _IdentityEntity(TypedDict):
-    """Identity entity with required id and optional name."""
-
-    id: str
-    name: NotRequired[str]
-
-
-class Identity(TypedDict, total=False):
-    """
-    Identity structure for user and organization tracking.
-
-    Example:
-        {
-            "organization": {"id": "123", "name": "ACME"},
-            "user": {"id": "456", "name": "John Doe"}
-        }
-    """
-
-    organization: NotRequired[_IdentityEntity]
-    user: NotRequired[_IdentityEntity]
-
-
 class StartObserve(ContextDecorator):
     """
     Entry point for Basalt observability.
@@ -822,22 +801,23 @@ def add_evaluators(*slugs: str) -> None:
             handle.add_evaluators(*slugs)
 
     @staticmethod
-    def set_identity(
-        *,
-        user_id: str | None = None,
-        user_name: str | None = None,
-        organization_id: str | None = None,
-        organization_name: str | None = None,
-    ) -> None:
-        """Set identity on the current span."""
+    def set_identity(identity: Identity | None = None) -> None:
+        """
+        Set identity on the current span.
+
+        Args:
+            identity: Identity TypedDict with optional 'user' and 'organization' keys.
+                Each key should contain a dict with 'id' (required) and 'name' (optional).
+
+        Example:
+            >>> Observe.set_identity({
+            ...     "user": {"id": "user-123", "name": "John Doe"},
+            ...     "organization": {"id": "org-456", "name": "ACME Corp"}
+            ... })
+        """
         handle = get_current_span_handle()
         if handle:
-            handle.set_identity(
-                user_id=user_id,
-                user_name=user_name,
-                organization_id=organization_id,
-                organization_name=organization_name,
-            )
+            handle.set_identity(identity)
 
 
 class AsyncStartObserve:

basalt/observability/context_managers.py

@@ -29,6 +29,7 @@
     apply_organization_from_context,
     apply_user_from_context,
 )
+from .types import Identity
 
 SPAN_TYPE_ATTRIBUTE = semconv.BasaltSpan.KIND
 EVALUATOR_CONTEXT_KEY: Final[str] = "basalt.context.evaluators"
@@ -364,31 +365,30 @@ def add_evaluators(self, *evaluators: Sequence[str]) -> None:
         for attachment in normalize_evaluator_specs(evaluators):
             self._append_evaluator(attachment)
 
-    def set_identity(
-        self,
-        *,
-        user_id: str | None = None,
-        user_name: str | None = None,
-        organization_id: str | None = None,
-        organization_name: str | None = None,
-    ) -> None:
+    def set_identity(self, identity: Identity | None = None) -> None:
         """
         Set user and/or organization identity for the span.
 
         Args:
-            user_id: User identifier to associate with the span.
-            user_name: Optional user display name.
-            organization_id: Organization identifier to associate with the span.
-            organization_name: Optional organization display name.
-        """
-        if user_id is not None:
-            self._span.set_attribute(semconv.BasaltUser.ID, user_id)
-            if user_name is not None:
-                self._span.set_attribute(semconv.BasaltUser.NAME, user_name)
-        if organization_id is not None:
-            self._span.set_attribute(semconv.BasaltOrganization.ID, organization_id)
-            if organization_name is not None:
-                self._span.set_attribute(semconv.BasaltOrganization.NAME, organization_name)
+            identity: Identity TypedDict with optional 'user' and 'organization' keys.
+                Each key should contain a dict with 'id' (required) and 'name' (optional).
+
+        Example:
+            >>> span.set_identity({
+            ...     "user": {"id": "user-123", "name": "John Doe"},
+            ...     "organization": {"id": "org-456", "name": "ACME Corp"}
+            ... })
+        """
+        if identity is None:
+            return
+
+        user_spec = identity.get("user") if identity else None
+        org_spec = identity.get("organization") if identity else None
+
+        if user_spec is not None:
+            apply_user_from_context(self._span, user_spec)
+        if org_spec is not None:
+            apply_organization_from_context(self._span, org_spec)
 
     def _append_evaluator(self, attachment: EvaluatorAttachment) -> None:
         """Attach an evaluator to the span, avoiding duplicates.
@@ -728,6 +728,27 @@ def _with_span_handle(
             # Mark all basalt spans with basalt.in_trace
             span.set_attribute(semconv.BasaltSpan.IN_TRACE, True)
 
+            # Inject prompt context if available
+            try:
+                from basalt.prompts.models import _current_prompt_context
+                prompt_ctx = _current_prompt_context.get()
+                if prompt_ctx:
+                    # Inject prompt attributes into this span
+                    import json
+                    span.set_attribute("basalt.prompt.slug", prompt_ctx["slug"])
+                    if prompt_ctx.get("version"):
+                        span.set_attribute("basalt.prompt.version", prompt_ctx["version"])
+                    if prompt_ctx.get("tag"):
+                        span.set_attribute("basalt.prompt.tag", prompt_ctx["tag"])
+                    span.set_attribute("basalt.prompt.model.provider", prompt_ctx["provider"])
+                    span.set_attribute("basalt.prompt.model.model", prompt_ctx["model"])
+                    if prompt_ctx.get("variables"):
+                        span.set_attribute("basalt.prompt.variables", json.dumps(prompt_ctx["variables"]))
+                    span.set_attribute("basalt.prompt.from_cache", prompt_ctx["from_cache"])
+            except ImportError:
+                # Prompts module not available, skip injection
+                pass
+
             _attach_attributes(span, attributes)
 
             # Apply metadata if provided
@@ -840,6 +861,27 @@ async def _async_with_span_handle(
             # Mark all basalt spans with basalt.in_trace
             span.set_attribute(semconv.BasaltSpan.IN_TRACE, True)
 
+            # Inject prompt context if available
+            try:
+                from basalt.prompts.models import _current_prompt_context
+                prompt_ctx = _current_prompt_context.get()
+                if prompt_ctx:
+                    # Inject prompt attributes into this span
+                    import json
+                    span.set_attribute("basalt.prompt.slug", prompt_ctx["slug"])
+                    if prompt_ctx.get("version"):
+                        span.set_attribute("basalt.prompt.version", prompt_ctx["version"])
+                    if prompt_ctx.get("tag"):
+                        span.set_attribute("basalt.prompt.tag", prompt_ctx["tag"])
+                    span.set_attribute("basalt.prompt.model.provider", prompt_ctx["provider"])
+                    span.set_attribute("basalt.prompt.model.model", prompt_ctx["model"])
+                    if prompt_ctx.get("variables"):
+                        span.set_attribute("basalt.prompt.variables", json.dumps(prompt_ctx["variables"]))
+                    span.set_attribute("basalt.prompt.from_cache", prompt_ctx["from_cache"])
+            except ImportError:
+                # Prompts module not available, skip injection
+                pass
+
             _attach_attributes(span, attributes)
 
             # Apply metadata if provided

basalt/observability/processors.py

@@ -328,6 +328,38 @@ def on_start(self, span: Span, parent_context: Any | None = None) -> None:  # ty
             # Apply prompt attributes (slug, version, provider, model)
             for key, value in prompt_data.items():
                 span.set_attribute(f"basalt.prompt.{key}", value)
+        else:
+            # If no explicit injection, try to read from ContextVar
+            # This allows auto-instrumented spans to inherit prompt context
+            # from the parent prompt context manager
+            try:
+                from basalt.prompts.models import _current_prompt_context
+                prompt_ctx = _current_prompt_context.get()
+                if prompt_ctx:
+                    import logging
+                    logger = logging.getLogger(__name__)
+                    logger.debug(f"✓ Injecting prompt context from ContextVar for span '{scope.name}': slug='{prompt_ctx['slug']}'")
+                    # Inject prompt attributes from ContextVar
+                    span.set_attribute("basalt.prompt.slug", prompt_ctx["slug"])
+                    if prompt_ctx.get("version"):
+                        span.set_attribute("basalt.prompt.version", prompt_ctx["version"])
+                    if prompt_ctx.get("tag"):
+                        span.set_attribute("basalt.prompt.tag", prompt_ctx["tag"])
+                    span.set_attribute("basalt.prompt.model.provider", prompt_ctx["provider"])
+                    span.set_attribute("basalt.prompt.model.model", prompt_ctx["model"])
+                    if prompt_ctx.get("variables"):
+                        span.set_attribute("basalt.prompt.variables", json.dumps(prompt_ctx["variables"]))
+                    span.set_attribute("basalt.prompt.from_cache", prompt_ctx["from_cache"])
+                else:
+                    import logging
+                    logger = logging.getLogger(__name__)
+                    logger.debug(f"✗ No prompt context found in ContextVar for auto-instrumented span '{scope.name}'")
+            except (ImportError, LookupError) as e:
+                import logging
+                logger = logging.getLogger(__name__)
+                logger.debug(f"✗ Failed to read prompt context for span '{scope.name}': {e}")
+                # Prompts module not available or no context set - skip injection
+                pass
 
         # Clear all pending injection data (single-use semantics)
         # We create a new context with all injection keys set to None

basalt/observability/request_tracing.py

@@ -51,8 +51,16 @@ async def trace_async_request(
             )
             raise
 
+        # Type-safe output formatting for PromptRequestSpan
+        from basalt.prompts.client import PromptRequestSpan
+        if isinstance(span_data, PromptRequestSpan):
+            output = span_data.format_output(result)
+        else:
+            status_code = getattr(result, "status_code", None)
+            output = {"status_code": status_code}
+
+        observe.set_output(output)
         status_code = getattr(result, "status_code", None)
-        observe.set_output({"status_code": status_code})
         span_data.finalize(
             span,
             duration_s=time.perf_counter() - start,
@@ -101,8 +109,16 @@ def trace_sync_request(
             )
             raise
 
+        # Type-safe output formatting for PromptRequestSpan
+        from basalt.prompts.client import PromptRequestSpan
+        if isinstance(span_data, PromptRequestSpan):
+            output = span_data.format_output(result)
+        else:
+            status_code = getattr(result, "status_code", None)
+            output = {"status_code": status_code}
+
+        observe.set_output(output)
         status_code = getattr(result, "status_code", None)
-        observe.set_output({"status_code": status_code})
         span_data.finalize(
             span,
             duration_s=time.perf_counter() - start,

basalt/observability/semconv.py

@@ -362,6 +362,13 @@ class BasaltAPI:
     Examples: "get", "list", "create", "update"
     """
 
+    INTERNAL: Final[str] = "basalt.internal.api"
+    """
+    Marks internal Basalt SDK API calls.
+    Type: boolean
+    Value: True for SDK->API requests (prompts, datasets, etc.)
+    """
+
 
 class BasaltCache:
     """Basalt caching attributes."""

basalt/observability/spans.py

@@ -40,6 +40,7 @@ def start_attributes(self) -> dict[str, Any]:
         attributes: dict[str, Any] = {
             semconv.BasaltAPI.CLIENT: self.client,
             semconv.BasaltAPI.OPERATION: self.operation,
+            semconv.BasaltAPI.INTERNAL: True,
             semconv.HTTP.METHOD: self.method.upper(),
             semconv.HTTP.URL: self.url,
         }

basalt/observability/types.py

@@ -0,0 +1,27 @@
+"""Shared observability types."""
+
+from __future__ import annotations
+
+from typing import TypedDict
+
+
+class _IdentityEntity(TypedDict):
+    """Identity entity with required id and optional name."""
+
+    id: str
+    name: str
+
+
+class Identity(TypedDict, total=False):
+    """
+    Identity structure for user and organization tracking.
+
+    Example:
+        {
+            "organization": {"id": "123", "name": "ACME"},
+            "user": {"id": "456", "name": "John Doe"}
+        }
+    """
+
+    organization: _IdentityEntity
+    user: _IdentityEntity