diff --git a/.env.example b/.env.example index ff244ef5..c9b87a7b 100644 --- a/.env.example +++ b/.env.example @@ -28,7 +28,7 @@ REDIS_URL=redis://localhost:6379/0 # --- REST API (Python FastAPI) ------------------------------------------------ # NOTE: Do NOT use generic HOST= or PORT= here — Next.js also reads PORT # and would start on the wrong port. FastAPI defaults to 0.0.0.0:8000. -CORS_ORIGINS=http://localhost:3000 +CORS_ORIGINS=["http://localhost:3000"] # --- Internal API (Python <-> Next.js service communication) ------------------ TRACEROOT_UI_URL=http://localhost:3000 diff --git a/backend/db/clickhouse/client.py b/backend/db/clickhouse/client.py index ea5c4e38..aec2aeb7 100644 --- a/backend/db/clickhouse/client.py +++ b/backend/db/clickhouse/client.py @@ -1,12 +1,13 @@ """ClickHouse client using clickhouse-connect.""" -import os from datetime import UTC, datetime from typing import Any import clickhouse_connect from clickhouse_connect.driver.client import Client +from shared.config import settings + class ClickHouseClient: """ClickHouse client wrapper for trace data operations.""" @@ -15,15 +16,15 @@ def __init__(self, client: Client): self._client = client @classmethod - def from_env(cls) -> "ClickHouseClient": - """Create client from environment variables.""" - port = os.getenv("CLICKHOUSE_HTTP_PORT") or os.getenv("CLICKHOUSE_PORT", "8123") + def from_settings(cls) -> "ClickHouseClient": + """Create client from centralized settings.""" + ch = settings.clickhouse client = clickhouse_connect.get_client( - host=os.getenv("CLICKHOUSE_HOST", "localhost"), - port=int(port), - username=os.getenv("CLICKHOUSE_USER", "clickhouse"), - password=os.getenv("CLICKHOUSE_PASSWORD", "clickhouse"), - database=os.getenv("CLICKHOUSE_DATABASE", "default"), + host=ch.host, + port=ch.port, + username=ch.user, + password=ch.password, + database=ch.database, ) return cls(client) @@ -148,5 +149,5 @@ def get_clickhouse_client() -> ClickHouseClient: """Get or create the singleton ClickHouse client.""" global _client if _client is None: - _client = ClickHouseClient.from_env() + _client = ClickHouseClient.from_settings() return _client diff --git a/backend/rest/main.py b/backend/rest/main.py index 489d92f2..a388e5fc 100644 --- a/backend/rest/main.py +++ b/backend/rest/main.py @@ -19,6 +19,7 @@ from rest.routers.public.traces import router as public_traces_router from rest.routers.traces import router as traces_router from rest.routers.users import router as users_router +from shared.config import settings app = FastAPI( title="Traceroot API", @@ -27,10 +28,9 @@ ) # CORS configuration -cors_origins = os.getenv("CORS_ORIGINS", "http://localhost:3000").split(",") app.add_middleware( CORSMiddleware, - allow_origins=cors_origins, + allow_origins=settings.cors_origins, allow_credentials=True, allow_methods=["*"], allow_headers=["*"], diff --git a/backend/rest/routers/deps.py b/backend/rest/routers/deps.py index e4e26a50..31cf9bae 100644 --- a/backend/rest/routers/deps.py +++ b/backend/rest/routers/deps.py @@ -1,15 +1,12 @@ """FastAPI dependencies for authentication (via Next.js internal API).""" -import os from typing import Annotated import httpx from fastapi import Depends, Header, HTTPException, status -from shared.enums import MemberRole -# Configuration for internal API (Python → Next.js) -TRACEROOT_UI_URL = os.getenv("TRACEROOT_UI_URL", "http://localhost:3000") -INTERNAL_API_SECRET = os.getenv("INTERNAL_API_SECRET", "") +from shared.config import settings +from shared.enums import MemberRole class ProjectAccessInfo: @@ -43,9 +40,9 @@ async def get_project_access( try: async with httpx.AsyncClient(timeout=10.0) as client: response = await client.post( - f"{TRACEROOT_UI_URL}/api/internal/validate-project-access", + f"{settings.traceroot_ui_url}/api/internal/validate-project-access", json={"userId": x_user_id, "projectId": project_id}, - headers={"X-Internal-Secret": INTERNAL_API_SECRET}, + headers={"X-Internal-Secret": settings.internal_api_secret}, ) except httpx.RequestError as e: raise HTTPException( diff --git a/backend/rest/routers/public/traces.py b/backend/rest/routers/public/traces.py index a3a2149f..7ad36883 100644 --- a/backend/rest/routers/public/traces.py +++ b/backend/rest/routers/public/traces.py @@ -14,7 +14,6 @@ import gzip import hashlib import logging -import os import uuid from datetime import UTC, datetime from typing import Annotated, Any @@ -28,16 +27,13 @@ from pydantic import BaseModel from rest.services.s3 import get_s3_service -from worker.tasks import process_s3_traces +from shared.config import settings +from worker.ingest_tasks import process_s3_traces logger = logging.getLogger(__name__) router = APIRouter(prefix="/public/traces", tags=["Traces (Public)"]) -# Configuration for internal API (Python → Next.js) -TRACEROOT_UI_URL = os.getenv("TRACEROOT_UI_URL", "http://localhost:3000") -INTERNAL_API_SECRET = os.getenv("INTERNAL_API_SECRET", "") - async def authenticate_api_key( authorization: Annotated[str | None, Header()] = None, @@ -81,9 +77,9 @@ async def authenticate_api_key( try: async with httpx.AsyncClient(timeout=10.0) as client: response = await client.post( - f"{TRACEROOT_UI_URL}/api/internal/validate-api-key", + f"{settings.traceroot_ui_url}/api/internal/validate-api-key", json={"keyHash": key_hash}, - headers={"X-Internal-Secret": INTERNAL_API_SECRET}, + headers={"X-Internal-Secret": settings.internal_api_secret}, ) except httpx.RequestError as e: logger.error(f"Failed to validate API key: {e}") diff --git a/backend/rest/services/s3.py b/backend/rest/services/s3.py index f033cc33..f4add10e 100644 --- a/backend/rest/services/s3.py +++ b/backend/rest/services/s3.py @@ -15,13 +15,14 @@ """ import logging -import os from typing import Any import boto3 from botocore.config import Config from botocore.exceptions import ClientError +from shared.config import settings + logger = logging.getLogger(__name__) @@ -39,17 +40,18 @@ def __init__( """Initialize S3 service. Args: - endpoint_url: S3/MinIO endpoint URL. - access_key_id: AWS access key ID. - secret_access_key: AWS secret access key. - bucket_name: S3 bucket name. - region: AWS region. + endpoint_url: S3/MinIO endpoint URL. Defaults to settings. + access_key_id: AWS access key ID. Defaults to settings. + secret_access_key: AWS secret access key. Defaults to settings. + bucket_name: S3 bucket name. Defaults to settings. + region: AWS region. Defaults to settings. """ - self._endpoint_url = endpoint_url or os.getenv("S3_ENDPOINT_URL") - self._access_key_id = access_key_id or os.getenv("S3_ACCESS_KEY_ID") - self._secret_access_key = secret_access_key or os.getenv("S3_SECRET_ACCESS_KEY") - self._bucket_name = bucket_name or os.getenv("S3_BUCKET_NAME", "traceroot") - self._region = region or os.getenv("S3_REGION", "us-east-1") + s3 = settings.s3 + self._endpoint_url = endpoint_url or s3.endpoint_url + self._access_key_id = access_key_id or s3.access_key_id + self._secret_access_key = secret_access_key or s3.secret_access_key + self._bucket_name = bucket_name or s3.bucket_name + self._region = region or s3.region self._client: Any = None diff --git a/backend/shared/config.py b/backend/shared/config.py new file mode 100644 index 00000000..ca46b568 --- /dev/null +++ b/backend/shared/config.py @@ -0,0 +1,78 @@ +"""Centralized configuration for the Traceroot backend. + +All environment variables are read once at import time via Pydantic Settings. +Services import ``settings`` from this module instead of calling os.getenv(). + +Environment variables are loaded from .env by entrypoints (rest/main.py, +worker/celery_app.py) before this module is first imported. +""" + +from pydantic_settings import BaseSettings, SettingsConfigDict + + +class ClickHouseSettings(BaseSettings): + """ClickHouse connection settings. + + Env vars: CLICKHOUSE_HOST, CLICKHOUSE_PORT, CLICKHOUSE_NATIVE_PORT, + CLICKHOUSE_USER, CLICKHOUSE_PASSWORD, CLICKHOUSE_DATABASE + """ + + model_config = SettingsConfigDict(env_prefix="CLICKHOUSE_") + + host: str = "localhost" + port: int = 8123 + native_port: int = 9000 + user: str = "clickhouse" + password: str = "clickhouse" + database: str = "default" + + +class S3Settings(BaseSettings): + """S3/MinIO settings for trace data storage. + + Env vars: S3_ENDPOINT_URL, S3_ACCESS_KEY_ID, S3_SECRET_ACCESS_KEY, + S3_BUCKET_NAME, S3_REGION + """ + + model_config = SettingsConfigDict(env_prefix="S3_") + + endpoint_url: str | None = None + access_key_id: str | None = None + secret_access_key: str | None = None + bucket_name: str = "traceroot" + region: str = "us-east-1" + + +class RedisSettings(BaseSettings): + """Redis settings for Celery broker and result backend. + + Env vars: REDIS_URL, REDIS_RESULT_URL + """ + + model_config = SettingsConfigDict(env_prefix="REDIS_") + + url: str = "redis://localhost:6379/0" + result_url: str = "redis://localhost:6379/1" + + +class Settings(BaseSettings): + """Root settings for the Traceroot backend. + + Nested settings (clickhouse, s3, redis) each read from their own + prefixed env vars. Top-level fields read from unprefixed env vars. + """ + + # CORS + cors_origins: list[str] = ["http://localhost:3000"] + + # Internal communication (Python <-> Next.js) + traceroot_ui_url: str = "http://localhost:3000" + internal_api_secret: str = "" + + # Service-specific settings + clickhouse: ClickHouseSettings = ClickHouseSettings() + s3: S3Settings = S3Settings() + redis: RedisSettings = RedisSettings() + + +settings = Settings() diff --git a/backend/worker/celery_app.py b/backend/worker/celery_app.py index bb0e8b71..2154cedd 100644 --- a/backend/worker/celery_app.py +++ b/backend/worker/celery_app.py @@ -18,22 +18,20 @@ # Load environment variables from .env file load_dotenv() +from shared.config import settings + # Configure logging for worker tasks logging.basicConfig( level=logging.INFO, format="[%(asctime)s: %(levelname)s/%(processName)s] %(name)s - %(message)s", ) -# Redis URL from environment or default -REDIS_URL = os.getenv("REDIS_URL", "redis://localhost:6379/0") -REDIS_RESULT_URL = os.getenv("REDIS_RESULT_URL", "redis://localhost:6379/1") - app = Celery("traceroot") app.conf.update( # Broker and backend - broker_url=REDIS_URL, - result_backend=REDIS_RESULT_URL, + broker_url=settings.redis.url, + result_backend=settings.redis.result_url, # Reliability settings task_acks_late=True, # ACK after task completes (not before) task_reject_on_worker_lost=True, # Requeue if worker dies mid-task @@ -53,8 +51,8 @@ result_expires=3600, ) -# Auto-discover tasks from worker.tasks module -app.autodiscover_tasks(["worker"]) +# Auto-discover tasks from worker.ingest_tasks module +app.autodiscover_tasks(["worker"], related_name="ingest_tasks") @worker_ready.connect @@ -62,6 +60,7 @@ def on_worker_ready(**kwargs): """Run ClickHouse migrations when worker starts.""" logger.info("Running ClickHouse migrations on worker startup...") try: + ch = settings.clickhouse result = subprocess.run( [ str(Path(__file__).resolve().parent.parent / "db" / "clickhouse" / "migrate.sh"), @@ -71,11 +70,11 @@ def on_worker_ready(**kwargs): text=True, env={ **os.environ, - "CLICKHOUSE_HOST": os.getenv("CLICKHOUSE_HOST", "localhost"), - "CLICKHOUSE_PORT": os.getenv("CLICKHOUSE_NATIVE_PORT", "9000"), - "CLICKHOUSE_USER": os.getenv("CLICKHOUSE_USER", "clickhouse"), - "CLICKHOUSE_PASSWORD": os.getenv("CLICKHOUSE_PASSWORD", "clickhouse"), - "CLICKHOUSE_DATABASE": os.getenv("CLICKHOUSE_DATABASE", "default"), + "CLICKHOUSE_HOST": ch.host, + "CLICKHOUSE_PORT": str(ch.native_port), + "CLICKHOUSE_USER": ch.user, + "CLICKHOUSE_PASSWORD": ch.password, + "CLICKHOUSE_DATABASE": ch.database, }, ) diff --git a/backend/worker/features/__init__.py b/backend/worker/features/__init__.py deleted file mode 100644 index 06faf2b0..00000000 --- a/backend/worker/features/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Worker features module.""" diff --git a/backend/worker/tasks.py b/backend/worker/ingest_tasks.py similarity index 97% rename from backend/worker/tasks.py rename to backend/worker/ingest_tasks.py index 1e60b7af..fd0fe10a 100644 --- a/backend/worker/tasks.py +++ b/backend/worker/ingest_tasks.py @@ -35,7 +35,7 @@ def process_s3_traces(self, s3_key: str, project_id: str) -> dict: # Import here to avoid circular imports and ensure fresh connections from db.clickhouse.client import get_clickhouse_client from rest.services.s3 import get_s3_service - from worker.transformer import transform_otel_to_clickhouse + from worker.otel_transform import transform_otel_to_clickhouse logger.info(f"Processing S3 traces: {s3_key} for project {project_id}") diff --git a/backend/worker/jobs/__init__.py b/backend/worker/jobs/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/backend/worker/jobs/s3_to_clickhouse.py b/backend/worker/jobs/s3_to_clickhouse.py deleted file mode 100644 index e69de29b..00000000 diff --git a/backend/worker/transformer.py b/backend/worker/otel_transform.py similarity index 99% rename from backend/worker/transformer.py rename to backend/worker/otel_transform.py index 6c052801..f460c519 100644 --- a/backend/worker/transformer.py +++ b/backend/worker/otel_transform.py @@ -332,7 +332,7 @@ def transform_otel_to_clickhouse( span_record["total_tokens"] = total_tokens # Calculate cost from actual token counts - from worker.features.tokens.pricing import get_model_price + from worker.tokens.pricing import get_model_price prices = get_model_price(model_name) if prices: @@ -351,7 +351,7 @@ def transform_otel_to_clickhouse( span_record["cost"] = float(input_cost + output_cost) else: # Fall back to text-based estimation - from worker.features.tokens import calculate_cost + from worker.tokens import calculate_cost usage = calculate_cost( model=model_name, diff --git a/backend/worker/processors/__init__.py b/backend/worker/processors/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/backend/worker/schedules/__init__.py b/backend/worker/schedules/__init__.py deleted file mode 100644 index 746a4ced..00000000 --- a/backend/worker/schedules/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Worker schedules.""" diff --git a/backend/worker/tests/__init__.py b/backend/worker/tests/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/backend/worker/tests/test_transformer.py b/backend/worker/tests/test_transformer.py deleted file mode 100644 index 080d405e..00000000 --- a/backend/worker/tests/test_transformer.py +++ /dev/null @@ -1,451 +0,0 @@ -"""Tests for OTEL → ClickHouse transformer. - -Covers: -- Traceroot SDK attributes (traceroot.span.input, traceroot.llm.model, etc.) -- OpenInference attributes (input.value, output.value, llm.token_count.*, etc.) -- GenAI semantic convention attributes (gen_ai.usage.*, gen_ai.request.model, etc.) -- Span kind detection from openinference.span.kind -- Fallback priority between attribute sources -""" - -import base64 -import json -import os -import sys - -import pytest - -# Add backend root to path so `worker.*` imports resolve -sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..")) - -from worker.transformer import transform_otel_to_clickhouse - -# ============================================================================= -# Helpers to build OTEL JSON payloads -# ============================================================================= - - -def _encode_id(hex_id: str) -> str: - """Encode a hex ID string to base64, matching OTLP wire format.""" - return base64.b64encode(bytes.fromhex(hex_id)).decode() - - -def _make_attr(key: str, value) -> dict: - """Build an OTEL attribute entry.""" - if isinstance(value, str): - return {"key": key, "value": {"stringValue": value}} - elif isinstance(value, bool): - return {"key": key, "value": {"boolValue": value}} - elif isinstance(value, int): - return {"key": key, "value": {"intValue": str(value)}} - elif isinstance(value, float): - return {"key": key, "value": {"doubleValue": value}} - else: - return {"key": key, "value": {"stringValue": str(value)}} - - -def _make_otel_payload(spans: list[dict]) -> dict: - """Wrap span dicts into a full OTEL resourceSpans payload.""" - return { - "resourceSpans": [ - { - "resource": {"attributes": []}, - "scopeSpans": [{"scope": {"name": "test"}, "spans": spans}], - } - ] - } - - -def _make_span( - trace_id: str = "0" * 32, - span_id: str = "0" * 16, - parent_span_id: str | None = None, - name: str = "test-span", - attributes: list[dict] | None = None, - kind: str = "SPAN_KIND_INTERNAL", -) -> dict: - """Build a single OTEL span dict.""" - span = { - "traceId": _encode_id(trace_id), - "spanId": _encode_id(span_id), - "name": name, - "kind": kind, - "startTimeUnixNano": "1700000000000000000", - "endTimeUnixNano": "1700000001000000000", - "attributes": attributes or [], - "status": {"code": 0}, - } - if parent_span_id: - span["parentSpanId"] = _encode_id(parent_span_id) - return span - - -# ============================================================================= -# Tests: Traceroot SDK attributes (baseline) -# ============================================================================= - - -class TestTracerootAttributes: - """Verify spans with traceroot.* attributes are processed correctly.""" - - def test_traceroot_input_output(self): - """traceroot.span.input and traceroot.span.output are extracted.""" - span = _make_span( - attributes=[ - _make_attr("traceroot.span.input", '{"prompt": "hello"}'), - _make_attr("traceroot.span.output", '{"response": "world"}'), - ] - ) - _, spans = transform_otel_to_clickhouse(_make_otel_payload([span]), "proj-1") - - assert len(spans) == 1 - assert spans[0]["input"] == '{"prompt": "hello"}' - assert spans[0]["output"] == '{"response": "world"}' - - def test_traceroot_llm_model(self): - """traceroot.llm.model sets model_name on LLM spans.""" - span = _make_span( - attributes=[ - _make_attr("traceroot.span.type", "llm"), - _make_attr("traceroot.llm.model", "gpt-4o-mini"), - _make_attr("traceroot.span.input", "test input"), - _make_attr("traceroot.span.output", "test output"), - ] - ) - _, spans = transform_otel_to_clickhouse(_make_otel_payload([span]), "proj-1") - - assert spans[0]["model_name"] == "gpt-4o-mini" - assert spans[0]["span_kind"] == "LLM" - - def test_traceroot_span_type_detection(self): - """traceroot.span.type correctly sets span_kind.""" - for span_type, expected_kind in [ - ("llm", "LLM"), - ("agent", "AGENT"), - ("tool", "TOOL"), - ("span", "SPAN"), - ]: - span = _make_span(attributes=[_make_attr("traceroot.span.type", span_type)]) - _, spans = transform_otel_to_clickhouse(_make_otel_payload([span]), "proj-1") - assert spans[0]["span_kind"] == expected_kind, f"type={span_type}" - - -# ============================================================================= -# Tests: OpenInference attributes (from OpenAIAgentsInstrumentor etc.) -# ============================================================================= - - -class TestOpenInferenceAttributes: - """Verify spans with OpenInference attributes are processed correctly. - - These attributes are set by openinference-instrumentation-openai-agents, - openinference-instrumentation-openai, and similar instrumentors. - """ - - def test_input_value_output_value(self): - """input.value and output.value are extracted as span input/output.""" - span = _make_span( - attributes=[ - _make_attr("input.value", '{"role": "user", "content": "What is AI?"}'), - _make_attr("output.value", '{"role": "assistant", "content": "AI is..."}'), - ] - ) - _, spans = transform_otel_to_clickhouse(_make_otel_payload([span]), "proj-1") - - assert spans[0]["input"] == '{"role": "user", "content": "What is AI?"}' - assert spans[0]["output"] == '{"role": "assistant", "content": "AI is..."}' - - def test_traceroot_attrs_take_priority_over_openinference(self): - """traceroot.span.input takes priority over input.value.""" - span = _make_span( - attributes=[ - _make_attr("traceroot.span.input", "traceroot input"), - _make_attr("input.value", "openinference input"), - _make_attr("traceroot.span.output", "traceroot output"), - _make_attr("output.value", "openinference output"), - ] - ) - _, spans = transform_otel_to_clickhouse(_make_otel_payload([span]), "proj-1") - - assert spans[0]["input"] == "traceroot input" - assert spans[0]["output"] == "traceroot output" - - def test_openinference_span_kind(self): - """openinference.span.kind maps to correct span_kind.""" - for oi_kind, expected in [ - ("LLM", "LLM"), - ("AGENT", "AGENT"), - ("TOOL", "TOOL"), - ("CHAIN", "SPAN"), - ]: - span = _make_span(attributes=[_make_attr("openinference.span.kind", oi_kind)]) - _, spans = transform_otel_to_clickhouse(_make_otel_payload([span]), "proj-1") - assert spans[0]["span_kind"] == expected, f"openinference.span.kind={oi_kind}" - - def test_llm_token_count_attributes(self): - """llm.token_count.* attributes provide accurate token usage.""" - span = _make_span( - attributes=[ - _make_attr("openinference.span.kind", "LLM"), - _make_attr("llm.model_name", "gpt-4o-mini"), - _make_attr("llm.token_count.prompt", 150), - _make_attr("llm.token_count.completion", 200), - _make_attr("llm.token_count.total", 350), - _make_attr("input.value", "some input"), - _make_attr("output.value", "some output"), - ] - ) - _, spans = transform_otel_to_clickhouse(_make_otel_payload([span]), "proj-1") - - assert spans[0]["input_tokens"] == 150 - assert spans[0]["output_tokens"] == 200 - assert spans[0]["total_tokens"] == 350 - - def test_token_counts_preferred_over_text_estimation(self): - """API token counts should be used instead of tiktoken estimation.""" - span = _make_span( - attributes=[ - _make_attr("openinference.span.kind", "LLM"), - _make_attr("llm.model_name", "gpt-4o-mini"), - _make_attr("llm.token_count.prompt", 42), - _make_attr("llm.token_count.completion", 73), - # input.value text would yield different token count if estimated - _make_attr("input.value", "a " * 500), - _make_attr("output.value", "b " * 500), - ] - ) - _, spans = transform_otel_to_clickhouse(_make_otel_payload([span]), "proj-1") - - # Should use API counts (42, 73), NOT text estimation - assert spans[0]["input_tokens"] == 42 - assert spans[0]["output_tokens"] == 73 - assert spans[0]["total_tokens"] == 42 + 73 - - def test_cost_calculated_from_api_token_counts(self): - """Cost should be calculated from API-provided token counts.""" - span = _make_span( - attributes=[ - _make_attr("openinference.span.kind", "LLM"), - _make_attr("llm.model_name", "gpt-4o-mini"), - _make_attr("llm.token_count.prompt", 1000000), # 1M input tokens - _make_attr("llm.token_count.completion", 1000000), # 1M output tokens - ] - ) - _, spans = transform_otel_to_clickhouse(_make_otel_payload([span]), "proj-1") - - # gpt-4o-mini: $0.15/1M input, $0.60/1M output - assert spans[0]["cost"] == pytest.approx(0.15 + 0.60, abs=0.001) - - def test_llm_model_name_fallback(self): - """llm.model_name is used when traceroot.llm.model is not set.""" - span = _make_span( - attributes=[ - _make_attr("openinference.span.kind", "LLM"), - _make_attr("llm.model_name", "gpt-4o"), - _make_attr("input.value", "hello"), - ] - ) - _, spans = transform_otel_to_clickhouse(_make_otel_payload([span]), "proj-1") - - assert spans[0]["model_name"] == "gpt-4o" - - def test_agent_span_with_model_and_tokens(self): - """AGENT spans with model/token attrs should still extract tokens. - - OpenAIAgentsInstrumentor sets openinference.span.kind=AGENT on - ResponseSpanData, but these spans carry llm.model_name and - llm.token_count.* from the API response. The transformer must - NOT gate token extraction on span_kind == "LLM". - """ - span = _make_span( - name="Optimist", - attributes=[ - _make_attr("openinference.span.kind", "AGENT"), # NOT "LLM" - _make_attr("llm.model_name", "gpt-4o-mini"), - _make_attr("llm.token_count.prompt", 100), - _make_attr("llm.token_count.completion", 80), - _make_attr("llm.token_count.total", 180), - _make_attr("input.value", "Give your perspective"), - _make_attr("output.value", "AI is great..."), - ], - ) - _, spans = transform_otel_to_clickhouse(_make_otel_payload([span]), "proj-1") - - s = spans[0] - assert s["span_kind"] == "AGENT" - assert s["model_name"] == "gpt-4o-mini" - assert s["input_tokens"] == 100 - assert s["output_tokens"] == 80 - assert s["total_tokens"] == 180 - assert s["cost"] is not None and s["cost"] > 0 - assert s["input"] == "Give your perspective" - assert s["output"] == "AI is great..." - - def test_chain_span_with_model_and_tokens(self): - """CHAIN/SPAN kind spans with model attrs should also extract tokens.""" - span = _make_span( - attributes=[ - _make_attr("openinference.span.kind", "CHAIN"), - _make_attr("llm.model_name", "gpt-4o"), - _make_attr("llm.token_count.prompt", 50), - _make_attr("llm.token_count.completion", 60), - ] - ) - _, spans = transform_otel_to_clickhouse(_make_otel_payload([span]), "proj-1") - - assert spans[0]["span_kind"] == "SPAN" - assert spans[0]["model_name"] == "gpt-4o" - assert spans[0]["input_tokens"] == 50 - assert spans[0]["output_tokens"] == 60 - - def test_full_openinference_llm_span(self): - """Full OpenInference LLM span with all attributes.""" - span = _make_span( - name="Optimist", - attributes=[ - _make_attr("openinference.span.kind", "LLM"), - _make_attr("llm.model_name", "gpt-4o-mini"), - _make_attr("llm.token_count.prompt", 85), - _make_attr("llm.token_count.completion", 120), - _make_attr("llm.token_count.total", 205), - _make_attr( - "input.value", - json.dumps([{"role": "user", "content": "Give your perspective on AI"}]), - ), - _make_attr( - "output.value", - json.dumps({"role": "assistant", "content": "AI is transformative..."}), - ), - ], - ) - _, spans = transform_otel_to_clickhouse(_make_otel_payload([span]), "proj-1") - - s = spans[0] - assert s["name"] == "Optimist" - assert s["span_kind"] == "LLM" - assert s["model_name"] == "gpt-4o-mini" - assert s["input_tokens"] == 85 - assert s["output_tokens"] == 120 - assert s["total_tokens"] == 205 - assert s["cost"] is not None and s["cost"] > 0 - assert "AI" in s["input"] - assert "transformative" in s["output"] - - -# ============================================================================= -# Tests: GenAI semantic convention attributes -# ============================================================================= - - -class TestGenAIAttributes: - """Verify spans with gen_ai.* attributes work correctly.""" - - def test_gen_ai_request_model(self): - """gen_ai.request.model is used as model_name fallback.""" - span = _make_span( - attributes=[ - _make_attr("openinference.span.kind", "LLM"), - _make_attr("gen_ai.request.model", "gpt-4o"), - _make_attr("input.value", "test"), - ] - ) - _, spans = transform_otel_to_clickhouse(_make_otel_payload([span]), "proj-1") - - assert spans[0]["model_name"] == "gpt-4o" - - def test_gen_ai_usage_token_counts(self): - """gen_ai.usage.* attributes provide token counts.""" - span = _make_span( - attributes=[ - _make_attr("openinference.span.kind", "LLM"), - _make_attr("llm.model_name", "gpt-4o-mini"), - _make_attr("gen_ai.usage.input_tokens", 100), - _make_attr("gen_ai.usage.output_tokens", 50), - ] - ) - _, spans = transform_otel_to_clickhouse(_make_otel_payload([span]), "proj-1") - - assert spans[0]["input_tokens"] == 100 - assert spans[0]["output_tokens"] == 50 - assert spans[0]["total_tokens"] == 150 - - def test_gen_ai_usage_prompt_completion_tokens(self): - """gen_ai.usage.prompt_tokens / completion_tokens also work.""" - span = _make_span( - attributes=[ - _make_attr("openinference.span.kind", "LLM"), - _make_attr("llm.model_name", "gpt-4o-mini"), - _make_attr("gen_ai.usage.prompt_tokens", 200), - _make_attr("gen_ai.usage.completion_tokens", 300), - _make_attr("gen_ai.usage.total_tokens", 500), - ] - ) - _, spans = transform_otel_to_clickhouse(_make_otel_payload([span]), "proj-1") - - assert spans[0]["input_tokens"] == 200 - assert spans[0]["output_tokens"] == 300 - assert spans[0]["total_tokens"] == 500 - - -# ============================================================================= -# Tests: Fallback to text estimation when no API token counts -# ============================================================================= - - -class TestTextEstimationFallback: - """Verify text-based token estimation works when API counts aren't available.""" - - def test_falls_back_to_text_estimation(self): - """Without API token counts, tokens are estimated from input/output text.""" - span = _make_span( - attributes=[ - _make_attr("traceroot.span.type", "llm"), - _make_attr("traceroot.llm.model", "gpt-4o-mini"), - _make_attr("traceroot.span.input", "Hello, how are you?"), - _make_attr("traceroot.span.output", "I am doing well, thank you!"), - ] - ) - _, spans = transform_otel_to_clickhouse(_make_otel_payload([span]), "proj-1") - - # Should have non-zero tokens from text estimation - assert spans[0]["input_tokens"] > 0 - assert spans[0]["output_tokens"] > 0 - - def test_no_input_output_yields_zero_tokens_on_fallback(self): - """With no input/output and no API counts, tokens should be zero.""" - span = _make_span( - attributes=[ - _make_attr("traceroot.span.type", "llm"), - _make_attr("traceroot.llm.model", "gpt-4o-mini"), - ] - ) - _, spans = transform_otel_to_clickhouse(_make_otel_payload([span]), "proj-1") - - assert spans[0]["input_tokens"] == 0 - assert spans[0]["output_tokens"] == 0 - - -# ============================================================================= -# Tests: Trace-level input/output from root span -# ============================================================================= - - -class TestTraceInputOutput: - """Verify trace-level records get input/output from root span.""" - - def test_root_span_openinference_input_output_propagates_to_trace(self): - """input.value/output.value on root span should set trace input/output.""" - root_span = _make_span( - trace_id="a" * 32, - span_id="1" * 16, - name="my-trace", - attributes=[ - _make_attr("input.value", "trace input text"), - _make_attr("output.value", "trace output text"), - ], - ) - traces, _ = transform_otel_to_clickhouse(_make_otel_payload([root_span]), "proj-1") - - assert len(traces) == 1 - assert traces[0]["input"] == "trace input text" - assert traces[0]["output"] == "trace output text" diff --git a/backend/worker/features/tokens/__init__.py b/backend/worker/tokens/__init__.py similarity index 100% rename from backend/worker/features/tokens/__init__.py rename to backend/worker/tokens/__init__.py diff --git a/backend/worker/features/tokens/pricing.py b/backend/worker/tokens/pricing.py similarity index 100% rename from backend/worker/features/tokens/pricing.py rename to backend/worker/tokens/pricing.py diff --git a/backend/worker/features/tokens/types.py b/backend/worker/tokens/types.py similarity index 100% rename from backend/worker/features/tokens/types.py rename to backend/worker/tokens/types.py diff --git a/backend/worker/features/tokens/usage.py b/backend/worker/tokens/usage.py similarity index 100% rename from backend/worker/features/tokens/usage.py rename to backend/worker/tokens/usage.py diff --git a/frontend/ui/src/app/support/page.tsx b/frontend/ui/src/app/support/page.tsx index af711b21..535d676f 100644 --- a/frontend/ui/src/app/support/page.tsx +++ b/frontend/ui/src/app/support/page.tsx @@ -3,6 +3,7 @@ import { useEffect } from "react"; import { Github, BookOpen, MessageCircle } from "lucide-react"; import { useLayout } from "@/components/layout/app-layout"; +import { clientEnv } from "@/env.client"; const supportChannels = [ { @@ -10,7 +11,7 @@ const supportChannels = [ title: "Documentation", description: "Tutorials and guides to get started.", icon: BookOpen, - href: process.env.NEXT_PUBLIC_DOCS_URL || "https://docs.traceroot.ai", + href: clientEnv.NEXT_PUBLIC_DOCS_URL, external: true, }, { @@ -18,9 +19,7 @@ const supportChannels = [ title: "GitHub Issues", description: "Report bugs or request new features.", icon: Github, - href: - process.env.NEXT_PUBLIC_GITHUB_ISSUES_URL || - "https://github.com/traceroot-ai/traceroot/issues", + href: clientEnv.NEXT_PUBLIC_GITHUB_ISSUES_URL, external: true, }, { @@ -28,7 +27,7 @@ const supportChannels = [ title: "Discord", description: "Chat with the community and team.", icon: MessageCircle, - href: process.env.NEXT_PUBLIC_DISCORD_INVITE_URL || "https://discord.com/invite/tPyffEZvvJ", + href: clientEnv.NEXT_PUBLIC_DISCORD_INVITE_URL, external: true, }, ]; diff --git a/frontend/ui/src/components/layout/sidebar.tsx b/frontend/ui/src/components/layout/sidebar.tsx index 51bcccde..3fffbf99 100644 --- a/frontend/ui/src/components/layout/sidebar.tsx +++ b/frontend/ui/src/components/layout/sidebar.tsx @@ -19,6 +19,7 @@ import { import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip"; import { Logo } from "@/components/Logo"; import { cn } from "@/lib/utils"; +import { clientEnv } from "@/env.client"; // Check if we're in a project context by looking at the path structure function getProjectContext(pathname: string): { isProject: boolean; projectId: string | null } { @@ -144,10 +145,7 @@ export function Sidebar({ collapsed = false }: SidebarProps) { { const { to, inviterName, inviterEmail, workspaceName, inviteId, role } = params; - const smtpUrl = process.env.TRACEROOT_SMTP_URL; - const mailFrom = process.env.TRACEROOT_SMTP_MAIL_FROM; - const baseUrl = process.env.NEXTAUTH_URL || "http://localhost:3000"; + const smtpUrl = env.TRACEROOT_SMTP_URL; + const mailFrom = env.TRACEROOT_SMTP_MAIL_FROM; + const baseUrl = env.NEXTAUTH_URL; if (!smtpUrl || !mailFrom) { console.warn( @@ -98,7 +99,7 @@ function buildHtmlEmail(params: EmailContentParams): string { - TraceRoot + TraceRoot diff --git a/pyproject.toml b/pyproject.toml index 54d8d72a..e929543a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,6 +7,7 @@ requires-python = ">=3.11" dependencies = [ # Core "pydantic>=2.0", + "pydantic-settings>=2.0", "cuid2>=2.0.0", # Database (PostgreSQL access moved to Prisma in ui/) "clickhouse-connect>=0.7", @@ -78,9 +79,10 @@ ignore = [ [tool.ruff.lint.per-file-ignores] "backend/rest/main.py" = ["E402"] # load_dotenv() must run before imports +"backend/worker/celery_app.py" = ["E402"] # load_dotenv() must run before imports [tool.ruff.lint.isort] -known-first-party = ["db", "rest", "worker", "traceroot"] +known-first-party = ["db", "rest", "shared", "worker", "traceroot"] [tool.ruff.format] quote-style = "double" diff --git a/uv.lock b/uv.lock index 2875dd94..837d49df 100644 --- a/uv.lock +++ b/uv.lock @@ -1145,6 +1145,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/36/c7/cfc8e811f061c841d7990b0201912c3556bfeb99cdcb7ed24adc8d6f8704/pydantic_core-2.41.5-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:56121965f7a4dc965bff783d70b907ddf3d57f6eba29b6d2e5dabfaf07799c51", size = 2145302, upload-time = "2025-11-04T13:43:46.64Z" }, ] +[[package]] +name = "pydantic-settings" +version = "2.12.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pydantic" }, + { name = "python-dotenv" }, + { name = "typing-inspection" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/43/4b/ac7e0aae12027748076d72a8764ff1c9d82ca75a7a52622e67ed3f765c54/pydantic_settings-2.12.0.tar.gz", hash = "sha256:005538ef951e3c2a68e1c08b292b5f2e71490def8589d4221b95dab00dafcfd0", size = 194184, upload-time = "2025-11-10T14:25:47.013Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c1/60/5d4751ba3f4a40a6891f24eec885f51afd78d208498268c734e256fb13c4/pydantic_settings-2.12.0-py3-none-any.whl", hash = "sha256:fddb9fd99a5b18da837b29710391e945b1e30c135477f484084ee513adb93809", size = 51880, upload-time = "2025-11-10T14:25:45.546Z" }, +] + [[package]] name = "pygments" version = "2.19.2" @@ -1685,6 +1699,7 @@ dependencies = [ { name = "opentelemetry-proto" }, { name = "protobuf" }, { name = "pydantic" }, + { name = "pydantic-settings" }, { name = "python-dotenv" }, { name = "python-multipart" }, { name = "redis" }, @@ -1713,6 +1728,7 @@ requires-dist = [ { name = "opentelemetry-proto", specifier = ">=1.20.0" }, { name = "protobuf", specifier = ">=4.0" }, { name = "pydantic", specifier = ">=2.0" }, + { name = "pydantic-settings", specifier = ">=2.0" }, { name = "pytest", marker = "extra == 'dev'", specifier = ">=7.0" }, { name = "pytest-asyncio", marker = "extra == 'dev'", specifier = ">=0.21" }, { name = "python-dotenv", specifier = ">=1.0" },