Merge pull request #256 from askui/fix/tracing

fix: clean up otel tracing and add workspace ID to trace resource

Author: philipph-askui

docs/09_observability_telemetry_tracing.md

@@ -92,7 +92,6 @@ Set the following environment variables to configure tracing. All variables use
 | `ASKUI__OTEL_ENDPOINT` | OTLP HTTP endpoint URL | — |
 | `ASKUI__OTEL_SERVICE_NAME` | Service name reported in traces | `askui-python-sdk` |
 | `ASKUI__OTEL_SERVICE_VERSION` | Service version reported in traces | Current package version |
-| `ASKUI__OTEL_CLUSTER_NAME` | Cluster name reported in traces | `askui-dev` |
 
 #### Linux & MacOS
 ```bash
@@ -132,7 +131,6 @@ from askui.telemetry.otel import OtelSettings
 
 settings = OtelSettings(
     enabled=True,
-    service_name="my-custom-service",
-    cluster_name="production",
+    service_name="my-custom-service"
 )
 ```

pdm.lock

@@ -228,7 +228,6 @@ otel = [
     "opentelemetry-api>=1.38.0",
     "opentelemetry-exporter-otlp-proto-http>=1.38.0",
     "opentelemetry-instrumentation-httpx>=0.59b0",
-    "opentelemetry-instrumentation-sqlalchemy>=0.59b0",
 ]
 vertex = [
     "google-cloud-aiplatform>=1.122.0",

pyproject.toml

@@ -291,12 +291,7 @@ def _execute_step(self) -> bool:
         self._switch_speaker_if_needed()
 
         # 2. Get next message(s) from speaker and add to history
-        logger.debug("Executing step with speaker: %s", self.current_speaker.name)
-        result: SpeakerResult = self.current_speaker.handle_step(
-            self, self.cache_manager
-        )
-        for message in result.messages_to_add:
-            self._add_message(message)
+        result = self._get_next_message()
 
         # 3. Execute tool calls if applicable
         continue_loop = False
@@ -318,6 +313,16 @@ def _execute_step(self) -> bool:
 
         return continue_loop
 
+    @tracer.start_as_current_span("_get_next_message")
+    def _get_next_message(self) -> SpeakerResult:
+        logger.debug("Executing step with speaker: %s", self.current_speaker.name)
+        result: SpeakerResult = self.current_speaker.handle_step(
+            self, self.cache_manager
+        )
+        for message in result.messages_to_add:
+            self._add_message(message)
+        return result
+
     @tracer.start_as_current_span("_execute_tools_if_present")
     def _execute_tools_if_present(self, message: MessageParam) -> MessageParam | None:
         """Execute tools if the message contains tool use blocks.

src/askui/models/shared/conversation.py

@@ -5,7 +5,7 @@
 from opentelemetry.sdk.resources import Resource
 from opentelemetry.sdk.trace import TracerProvider
 from opentelemetry.sdk.trace.export import BatchSpanProcessor
-from pydantic import Field, SecretStr, model_validator
+from pydantic import AliasChoices, Field, SecretStr, model_validator
 from pydantic_settings import BaseSettings, SettingsConfigDict
 from typing_extensions import Self
 
@@ -21,6 +21,7 @@ class OtelSettings(BaseSettings):
         env_prefix="ASKUI__OTEL_",
         case_sensitive=False,
         extra="ignore",
+        validate_by_name=True,
     )
 
     enabled: bool = Field(default=False)
@@ -38,7 +39,11 @@ class OtelSettings(BaseSettings):
         default=None,
         description="OTLP endpoint URL.",
     )
-    cluster_name: str = Field(default="askui-dev")
+    workspace_id: str | None = Field(
+        default=None,
+        validation_alias=AliasChoices("ASKUI_WORKSPACE_ID", "askui_workspace_id"),
+        description="AskUI workspace ID. Read from ASKUI_WORKSPACE_ID env var.",
+    )
 
     @model_validator(mode="after")
     def validate_secret_when_enabled(self) -> Self:
@@ -71,23 +76,20 @@ def setup_opentelemetry_tracing(settings: OtelSettings) -> None:
         from opentelemetry.instrumentation.httpx import (
             HTTPXClientInstrumentor,
         )
-        from opentelemetry.instrumentation.sqlalchemy import (
-            SQLAlchemyInstrumentor,
-        )
     except ImportError:
         logger.exception("Failed to set up OpenTelemetry Tracing.")
         return
 
     if not settings.enabled:
         return
 
-    resource = Resource.create(
-        {
-            "service.name": settings.service_name,
-            "service.version": settings.service_version,
-            "cluster.name": settings.cluster_name,
-        }
-    )
+    resource_attributes: dict[str, str] = {
+        "service.name": settings.service_name,
+        "service.version": settings.service_version,
+    }
+    if settings.workspace_id:
+        resource_attributes["askui.workspace.id"] = settings.workspace_id
+    resource = Resource.create(resource_attributes)
     provider = TracerProvider(resource=resource)
 
     base: str = settings.user + ":" + settings.secret.get_secret_value()  # type: ignore
@@ -104,4 +106,3 @@ def setup_opentelemetry_tracing(settings: OtelSettings) -> None:
     trace.set_tracer_provider(provider)
 
     HTTPXClientInstrumentor().instrument()
-    SQLAlchemyInstrumentor().instrument()