feat(ai-service): add OpenTelemetry tracing around LLM calls

app/ai-service/main.py

@@ -264,10 +264,15 @@ async def _log_rejection(
 ]
 
 
+from tracing.otel_setup import setup_tracing
+
+
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     logger.info("Starting up ChainForge AI Service...")
+    setup_tracing()
     if not settings.validate_api_keys():
+
         logger.warning("No API keys configured. AI features will be unavailable.")
     else:
         provider = settings.get_active_provider()

app/ai-service/requirements-prod.txt

@@ -49,3 +49,7 @@ slowapi==0.1.9
 # System monitoring
 psutil==5.9.8
 prometheus-client==0.20.0
+opentelemetry-api==1.25.0
+opentelemetry-sdk==1.25.0
+opentelemetry-exporter-otlp-proto-http==1.25.0
+

app/ai-service/requirements.txt

@@ -51,3 +51,7 @@ pytest-cov==4.1.0
 # Monitoring and Metrics
 prometheus-client==0.20.0
 psutil==5.9.8
+opentelemetry-api==1.25.0
+opentelemetry-sdk==1.25.0
+opentelemetry-exporter-otlp-proto-http==1.25.0
+

app/ai-service/services/humanitarian_verification.py

@@ -88,6 +88,7 @@ def verify_claim(
                             system_prompt=prompt["system"],
                             user_prompt=prompt["user"],
                             timeout=timeout,
+                            prompt_variant=prompt_variant,
                         )
                         parsed = self._parse_json_response(raw_content)
                         if breaker:
@@ -143,50 +144,72 @@ def _call_provider(
         system_prompt: str,
         user_prompt: str,
         timeout: Optional[float] = None,
+        prompt_variant: Optional[str] = None,
     ) -> str:
-        if provider == "test":
-            return self._call_test(model, system_prompt, user_prompt)
-        if provider == "openai":
-            return self._call_openai(model, system_prompt, user_prompt, timeout)
-        if provider == "groq":
-            return self._call_groq(model, system_prompt, user_prompt, timeout)
-        raise ValueError(f"Unsupported provider: {provider}")
+        from tracing.otel_setup import get_tracer
+        tracer = get_tracer()
+        with tracer.start_as_current_span("humanitarian_verification.call_provider") as span:
+            span.set_attribute("model", model)
+            if prompt_variant:
+                span.set_attribute("prompt_variant", prompt_variant)
+
+            if provider == "test":
+                return self._call_test(model, system_prompt, user_prompt)
+            if provider == "openai":
+                return self._call_openai(model, system_prompt, user_prompt, timeout, prompt_variant=prompt_variant)
+            if provider == "groq":
+                return self._call_groq(model, system_prompt, user_prompt, timeout, prompt_variant=prompt_variant)
+            raise ValueError(f"Unsupported provider: {provider}")
 
     def _call_openai(
         self,
         model: str,
         system_prompt: str,
         user_prompt: str,
         timeout: Optional[float] = None,
+        prompt_variant: Optional[str] = None,
     ) -> str:
-        if not settings.openai_api_key:
-            raise RuntimeError("OpenAI API key is not configured")
-        return self._call_chat_completion_api(
-            base_url="https://api.openai.com/v1/chat/completions",
-            api_key=settings.openai_api_key,
-            model=model,
-            system_prompt=system_prompt,
-            user_prompt=user_prompt,
-            timeout=timeout,
-        )
+        from tracing.otel_setup import get_tracer
+        tracer = get_tracer()
+        with tracer.start_as_current_span("humanitarian_verification.call_openai") as span:
+            span.set_attribute("model", model)
+            if prompt_variant:
+                span.set_attribute("prompt_variant", prompt_variant)
+            if not settings.openai_api_key:
+                raise RuntimeError("OpenAI API key is not configured")
+            return self._call_chat_completion_api(
+                base_url="https://api.openai.com/v1/chat/completions",
+                api_key=settings.openai_api_key,
+                model=model,
+                system_prompt=system_prompt,
+                user_prompt=user_prompt,
+                timeout=timeout,
+            )
 
     def _call_groq(
         self,
         model: str,
         system_prompt: str,
         user_prompt: str,
         timeout: Optional[float] = None,
+        prompt_variant: Optional[str] = None,
     ) -> str:
-        if not settings.groq_api_key:
-            raise RuntimeError("Groq API key is not configured")
-        return self._call_chat_completion_api(
-            base_url="https://api.groq.com/openai/v1/chat/completions",
-            api_key=settings.groq_api_key,
-            model=model,
-            system_prompt=system_prompt,
-            user_prompt=user_prompt,
-            timeout=timeout,
-        )
+        from tracing.otel_setup import get_tracer
+        tracer = get_tracer()
+        with tracer.start_as_current_span("humanitarian_verification.call_groq") as span:
+            span.set_attribute("model", model)
+            if prompt_variant:
+                span.set_attribute("prompt_variant", prompt_variant)
+            if not settings.groq_api_key:
+                raise RuntimeError("Groq API key is not configured")
+            return self._call_chat_completion_api(
+                base_url="https://api.groq.com/openai/v1/chat/completions",
+                api_key=settings.groq_api_key,
+                model=model,
+                system_prompt=system_prompt,
+                user_prompt=user_prompt,
+                timeout=timeout,
+            )
 
     def _call_chat_completion_api(
         self,

app/ai-service/tasks.py

@@ -6,10 +6,12 @@
 import logging
 import uuid
 import time
+import json
 from typing import Any, Dict, Optional
 from celery import Celery
 from celery.result import AsyncResult
 import httpx
+import redis
 
 import metrics
 from config import settings
@@ -70,8 +72,32 @@ def process_heavy_inference_task(self, task_id: str, payload: Dict[str, Any]) ->
 
     return process_heavy_inference_task
 
-# Task status storage (in production, use Redis with proper TTL)
-task_results: Dict[str, Dict[str, Any]] = {}
+# Lazy Redis client initialization - defers connection until needed
+redis_client = None
+
+def get_redis_client() -> redis.Redis:
+    """
+    Get or initialize the Redis client.
+    """
+    global redis_client
+    if redis_client is None:
+        redis_client = redis.from_url(settings.redis_url, decode_responses=True)
+    return redis_client
+
+
+def set_status(task_id: str, payload: Dict[str, Any]) -> None:
+    """
+    Write task status payload directly to Redis with a 24-hour TTL.
+    """
+    try:
+        r = get_redis_client()
+        key = f"task_status:{task_id}"
+        # TTL of 24 hours (86400 seconds)
+        r.setex(key, 86400, json.dumps(payload))
+    except Exception as e:
+        logger.error(f"Failed to write task status to Redis: {e}")
+
+
 pii_scrubber_service = PIIScrubberService()
 humanitarian_verification_service = HumanitarianVerificationService()
 
@@ -91,12 +117,13 @@ def update_task_status(
         result: Task result data (if completed)
         error: Error message (if failed)
     """
-    task_results[task_id] = {
+    payload = {
         'status': status,
         'result': result,
         'error': error,
         'updated_at': time.time()
     }
+    set_status(task_id, payload)
 
 
 def send_webhook_notification(task_id: str, status: str, result: Any = None, error: str = None) -> None:
@@ -373,20 +400,22 @@ def get_task_status(task_id: str) -> Dict[str, Any]:
                 'task_id': task_id,
                 'status': 'processing',
             }
-        else:
-            return {
-                'task_id': task_id,
-                'status': 'pending',
-            }
     except Exception:
         pass
 
-    # Fall back to local storage
-    if task_id in task_results:
-        return {
-            'task_id': task_id,
-            **task_results[task_id]
-        }
+    # Fall back to Redis storage
+    try:
+        r = get_redis_client()
+        key = f"task_status:{task_id}"
+        data = r.get(key)
+        if data:
+            payload = json.loads(data)
+            return {
+                'task_id': task_id,
+                **payload
+            }
+    except Exception as e:
+        logger.error(f"Failed to read task status from Redis: {e}")
 
     return {
         'task_id': task_id,

app/ai-service/tests/test_otel_spans.py

@@ -0,0 +1,106 @@
+import os
+import pytest
+from unittest.mock import patch, MagicMock
+from config import settings
+from services.humanitarian_verification import HumanitarianVerificationService
+from tracing.otel_setup import (
+    reset_tracing_for_test,
+    get_in_memory_exporter,
+)
+
+class TestOtelSpans:
+    @pytest.fixture(autouse=True)
+    def setup_otel(self, monkeypatch):
+        # Force app env to test to register InMemorySpanExporter
+        monkeypatch.setenv("APP_ENV", "test")
+        monkeypatch.setattr(settings, "openai_api_key", "test-openai-key")
+        monkeypatch.setattr(settings, "groq_api_key", "test-groq-key")
+        reset_tracing_for_test()
+
+        self.exporter = get_in_memory_exporter()
+        if self.exporter:
+            self.exporter.clear()
+
+    def test_verify_claim_emits_two_spans_openai(self, monkeypatch):
+        service = HumanitarianVerificationService()
+
+        # Mock _call_chat_completion_api to avoid making real network requests
+        mock_response = '{"verdict": "credible", "confidence": 0.95, "summary": "verified"}'
+        monkeypatch.setattr(
+            service,
+            "_call_chat_completion_api",
+            lambda *args, **kwargs: mock_response
+        )
+
+        # Ensure we only try to call openai
+        monkeypatch.setattr(service, "_provider_attempt_order", lambda pref: ["openai"])
+        monkeypatch.setattr(service, "_get_model_for_provider", lambda prov: "gpt-4-test")
+
+        # Trigger claim verification
+        result = service.verify_claim(
+            aid_claim="Food packs delivered to flood zone.",
+            supporting_evidence=["waybill-102"],
+            context_factors={"weather": "clear"},
+            provider_preference="openai"
+        )
+
+        assert result["provider"] == "openai"
+        assert result["prompt_variant"] == "primary"
+
+        # Verify tracing spans
+        finished_spans = self.exporter.get_finished_spans()
+        assert len(finished_spans) == 2, f"Expected 2 spans, got {len(finished_spans)}"
+
+        # Spans are emitted as they finish:
+        # call_openai is nested inside call_provider, so call_openai finishes first!
+        span_openai = finished_spans[0]
+        span_provider = finished_spans[1]
+
+        assert span_openai.name == "humanitarian_verification.call_openai"
+        assert span_openai.attributes.get("model") == "gpt-4-test"
+        assert span_openai.attributes.get("prompt_variant") == "primary"
+
+        assert span_provider.name == "humanitarian_verification.call_provider"
+        assert span_provider.attributes.get("model") == "gpt-4-test"
+        assert span_provider.attributes.get("prompt_variant") == "primary"
+
+    def test_verify_claim_emits_two_spans_groq(self, monkeypatch):
+        service = HumanitarianVerificationService()
+
+        # Mock _call_chat_completion_api to avoid making real network requests
+        mock_response = '{"verdict": "not_credible", "confidence": 0.85, "summary": "no evidence"}'
+        monkeypatch.setattr(
+            service,
+            "_call_chat_completion_api",
+            lambda *args, **kwargs: mock_response
+        )
+
+        # Ensure we only try to call groq
+        monkeypatch.setattr(service, "_provider_attempt_order", lambda pref: ["groq"])
+        monkeypatch.setattr(service, "_get_model_for_provider", lambda prov: "llama3-groq-test")
+
+        # Trigger claim verification
+        result = service.verify_claim(
+            aid_claim="Medicines delivered to shelter.",
+            supporting_evidence=["receipt-44"],
+            context_factors={"region": "north"},
+            provider_preference="groq"
+        )
+
+        assert result["provider"] == "groq"
+        assert result["prompt_variant"] == "primary"
+
+        # Verify tracing spans
+        finished_spans = self.exporter.get_finished_spans()
+        assert len(finished_spans) == 2, f"Expected 2 spans, got {len(finished_spans)}"
+
+        span_groq = finished_spans[0]
+        span_provider = finished_spans[1]
+
+        assert span_groq.name == "humanitarian_verification.call_groq"
+        assert span_groq.attributes.get("model") == "llama3-groq-test"
+        assert span_groq.attributes.get("prompt_variant") == "primary"
+
+        assert span_provider.name == "humanitarian_verification.call_provider"
+        assert span_provider.attributes.get("model") == "llama3-groq-test"
+        assert span_provider.attributes.get("prompt_variant") == "primary"