aiqa-client 0.1.0__py3-none-any.whl → 0.1.2__py3-none-any.whl

aiqa/__init__.py

@@ -1,29 +1,66 @@
 """
 Python client for AIQA server - OpenTelemetry tracing decorators.
+
+IMPORTANT: Before using any AIQA functionality, you must call get_aiqa_client() to initialize
+the client and load environment variables (AIQA_SERVER_URL, AIQA_API_KEY, AIQA_COMPONENT_TAG, etc.).
+
+Example:
+    from dotenv import load_dotenv
+    from aiqa import get_aiqa_client, WithTracing
+    
+    # Load environment variables from .env file (if using one)
+    load_dotenv()
+    
+    # Initialize client (must be called before using WithTracing or other functions)
+    get_aiqa_client()
+    
+    @WithTracing
+    def my_function():
+        return "Hello, AIQA!"
 """
 
 from .tracing import (
     WithTracing,
-    flush_spans,
+    flush_tracing,
     shutdown_tracing,
     set_span_attribute,
     set_span_name,
     get_active_span,
-    provider,
-    exporter,
+    get_provider,
+    get_exporter,
+    get_trace_id,
+    get_span_id,
+    create_span_from_trace_id,
+    inject_trace_context,
+    extract_trace_context,
+    set_conversation_id,
+    set_component_tag,
+    get_span,
 )
+from .client import get_aiqa_client
+from .experiment_runner import ExperimentRunner
 
-__version__ = "0.1.0"
+__version__ = "0.1.2"
 
 __all__ = [
     "WithTracing",
-    "flush_spans",
+    "flush_tracing",
     "shutdown_tracing",
     "set_span_attribute",
     "set_span_name",
     "get_active_span",
-    "provider",
-    "exporter",
+    "get_provider",
+    "get_exporter",
+    "get_aiqa_client",
+    "ExperimentRunner",
+    "get_trace_id",
+    "get_span_id",
+    "create_span_from_trace_id",
+    "inject_trace_context",
+    "extract_trace_context",
+    "set_conversation_id",
+    "set_component_tag",
+    "get_span",
     "__version__",
 ]
 

aiqa/aiqa_exporter.py

@@ -8,11 +8,12 @@ import json
 import logging
 import threading
 import time
+import io
 from typing import List, Dict, Any, Optional
 from opentelemetry.sdk.trace import ReadableSpan
 from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult
 
-logger = logging.getLogger(__name__)
+logger = logging.getLogger("AIQA")
 
 
 class AIQASpanExporter(SpanExporter):
@@ -39,10 +40,16 @@ class AIQASpanExporter(SpanExporter):
         self._api_key = api_key
         self.flush_interval_ms = flush_interval_seconds * 1000
         self.buffer: List[Dict[str, Any]] = []
+        self.buffer_span_keys: set = set()  # Track (traceId, spanId) tuples to prevent duplicates (Python 3.8 compatible)
         self.buffer_lock = threading.Lock()
         self.flush_lock = threading.Lock()
         self.shutdown_requested = False
         self.flush_timer: Optional[threading.Thread] = None
+        
+        logger.info(
+            f"Initializing AIQASpanExporter: server_url={self.server_url or 'not set'}, "
+            f"flush_interval={flush_interval_seconds}s"
+        )
         self._start_auto_flush()
 
     @property
@@ -56,14 +63,39 @@ class AIQASpanExporter(SpanExporter):
     def export(self, spans: List[ReadableSpan]) -> SpanExportResult:
         """
         Export spans to the AIQA server. Adds spans to buffer for async flushing.
+        Deduplicates spans based on (traceId, spanId) to prevent repeated exports.
         """
         if not spans:
+            logger.debug("export() called with empty spans list")
             return SpanExportResult.SUCCESS
-
-        # Serialize and add to buffer
+        logger.debug(f"AIQA export() called with {len(spans)} spans")
+        # Serialize and add to buffer, deduplicating by (traceId, spanId)
         with self.buffer_lock:
-            serialized_spans = [self._serialize_span(span) for span in spans]
+            serialized_spans = []
+            duplicates_count = 0
+            for span in spans:
+                serialized = self._serialize_span(span)
+                span_key = (serialized["traceId"], serialized["spanId"])
+                if span_key not in self.buffer_span_keys:
+                    serialized_spans.append(serialized)
+                    self.buffer_span_keys.add(span_key)
+                else:
+                    duplicates_count += 1
+                    logger.debug(f"export() skipping duplicate span: traceId={serialized['traceId']}, spanId={serialized['spanId']}")
+            
             self.buffer.extend(serialized_spans)
+            buffer_size = len(self.buffer)
+        
+        if duplicates_count > 0:
+            logger.debug(
+                f"export() added {len(serialized_spans)} span(s) to buffer, skipped {duplicates_count} duplicate(s). "
+                f"Total buffered: {buffer_size}"
+            )
+        else:
+            logger.debug(
+                f"export() added {len(spans)} span(s) to buffer. "
+                f"Total buffered: {buffer_size}"
+            )
 
         return SpanExportResult.SUCCESS
 
@@ -126,8 +158,8 @@ class AIQASpanExporter(SpanExporter):
             "duration": self._time_to_tuple(span.end_time - span.start_time) if span.end_time else None,
             "ended": span.end_time is not None,
             "instrumentationLibrary": {
-                "name": span.instrumentation_info.name if hasattr(span, "instrumentation_info") else "",
-                "version": span.instrumentation_info.version if hasattr(span, "instrumentation_info") else None,
+                "name": self._get_instrumentation_name(),
+                "version": self._get_instrumentation_version(),
             },
         }
 
@@ -136,18 +168,96 @@ class AIQASpanExporter(SpanExporter):
         seconds = int(nanoseconds // 1_000_000_000)
         nanos = int(nanoseconds % 1_000_000_000)
         return (seconds, nanos)
+    
+    def _get_instrumentation_name(self) -> str:
+        """Get instrumentation library name - always 'aiqa-tracer'."""
+        from .client import AIQA_TRACER_NAME
+        return AIQA_TRACER_NAME
+    
+    def _get_instrumentation_version(self) -> Optional[str]:
+        """Get instrumentation library version from __version__."""
+        try:
+            from . import __version__
+            return __version__
+        except (ImportError, AttributeError):
+            return None
+
+    def _build_request_headers(self) -> Dict[str, str]:
+        """Build HTTP headers for span requests."""
+        headers = {"Content-Type": "application/json"}
+        if self.api_key:
+            headers["Authorization"] = f"ApiKey {self.api_key}"
+        return headers
+
+    def _get_span_url(self) -> str:
+        """Get the URL for sending spans."""
+        if not self.server_url:
+            raise ValueError("AIQA_SERVER_URL is not set. Cannot send spans to server.")
+        return f"{self.server_url}/span"
+
+    def _is_interpreter_shutdown_error(self, error: Exception) -> bool:
+        """Check if error is due to interpreter shutdown."""
+        error_str = str(error)
+        return "cannot schedule new futures after" in error_str or "interpreter shutdown" in error_str
+
+    def _extract_spans_from_buffer(self) -> List[Dict[str, Any]]:
+        """Extract spans from buffer (thread-safe). Returns copy of buffer."""
+        with self.buffer_lock:
+            return self.buffer[:]
+
+    def _extract_and_remove_spans_from_buffer(self) -> List[Dict[str, Any]]:
+        """
+        Atomically extract and remove all spans from buffer (thread-safe).
+        Returns the extracted spans. This prevents race conditions where spans
+        are added between extraction and clearing.
+        Note: Does NOT clear buffer_span_keys - that should be done after successful send
+        to avoid unnecessary clearing/rebuilding on failures.
+        """
+        with self.buffer_lock:
+            spans = self.buffer[:]
+            self.buffer.clear()
+            return spans
+    
+    def _remove_span_keys_from_tracking(self, spans: List[Dict[str, Any]]) -> None:
+        """
+        Remove span keys from tracking set (thread-safe). Called after successful send.
+        """
+        with self.buffer_lock:
+            for span in spans:
+                span_key = (span["traceId"], span["spanId"])
+                self.buffer_span_keys.discard(span_key)
+
+    def _prepend_spans_to_buffer(self, spans: List[Dict[str, Any]]) -> None:
+        """
+        Prepend spans back to buffer (thread-safe). Used to restore spans
+        if sending fails. Rebuilds the span keys tracking set.
+        """
+        with self.buffer_lock:
+            self.buffer[:0] = spans
+            # Rebuild span keys set from current buffer contents
+            self.buffer_span_keys = {(span["traceId"], span["spanId"]) for span in self.buffer}
+
+    def _clear_buffer(self) -> None:
+        """Clear the buffer (thread-safe)."""
+        with self.buffer_lock:
+            self.buffer.clear()
+            self.buffer_span_keys.clear()
 
     async def flush(self) -> None:
         """
         Flush buffered spans to the server. Thread-safe: ensures only one flush operation runs at a time.
+        Atomically extracts spans to prevent race conditions with concurrent export() calls.
         """
+        logger.debug("flush() called - attempting to acquire flush lock")
         with self.flush_lock:
-            # Get current buffer and clear it atomically
-            with self.buffer_lock:
-                spans_to_flush = self.buffer[:]
-                self.buffer.clear()
+            logger.debug("flush() acquired flush lock")
+            # Atomically extract and remove spans to prevent race conditions
+            # where export() adds spans between extraction and clearing
+            spans_to_flush = self._extract_and_remove_spans_from_buffer()
+            logger.debug(f"flush() extracted {len(spans_to_flush)} span(s) from buffer")
 
             if not spans_to_flush:
+                logger.debug("flush() completed: no spans to flush")
                 return
 
             # Skip sending if server URL is not configured
@@ -155,94 +265,217 @@ class AIQASpanExporter(SpanExporter):
                 logger.warning(
                     f"Skipping flush: AIQA_SERVER_URL is not set. {len(spans_to_flush)} span(s) will not be sent."
                 )
+                # Spans already removed from buffer, clear their keys to free memory
+                self._remove_span_keys_from_tracking(spans_to_flush)
                 return
 
+            logger.info(f"flush() sending {len(spans_to_flush)} span(s) to server")
             try:
                 await self._send_spans(spans_to_flush)
+                logger.info(f"flush() successfully sent {len(spans_to_flush)} span(s) to server")
+                # Spans already removed from buffer during extraction
+                # Now clear their keys from tracking set to free memory
+                self._remove_span_keys_from_tracking(spans_to_flush)
+            except RuntimeError as error:
+                if self._is_interpreter_shutdown_error(error):
+                    if self.shutdown_requested:
+                        logger.debug(f"flush() skipped due to interpreter shutdown: {error}")
+                        # Put spans back for retry with sync send during shutdown
+                        self._prepend_spans_to_buffer(spans_to_flush)
+                    else:
+                        logger.warning(f"flush() interrupted by interpreter shutdown: {error}")
+                        # Put spans back for retry
+                        self._prepend_spans_to_buffer(spans_to_flush)
+                    raise
+                logger.error(f"Error flushing spans to server: {error}")
+                # Put spans back for retry
+                self._prepend_spans_to_buffer(spans_to_flush)
+                raise
             except Exception as error:
                 logger.error(f"Error flushing spans to server: {error}")
+                # Put spans back for retry
+                self._prepend_spans_to_buffer(spans_to_flush)
                 if self.shutdown_requested:
                     raise
 
     def _start_auto_flush(self) -> None:
         """Start the auto-flush timer."""
         if self.shutdown_requested:
+            logger.warning("_start_auto_flush() called but shutdown already requested")
             return
 
+        logger.info(f"Starting auto-flush thread with interval {self.flush_interval_ms / 1000.0}s")
+
         def flush_worker():
             import asyncio
+            logger.debug("Auto-flush worker thread started")
             loop = asyncio.new_event_loop()
             asyncio.set_event_loop(loop)
             
+            cycle_count = 0
             while not self.shutdown_requested:
+                cycle_count += 1
+                logger.debug(f"Auto-flush cycle #{cycle_count} starting")
                 try:
                     loop.run_until_complete(self.flush())
+                    logger.debug(f"Auto-flush cycle #{cycle_count} completed, sleeping {self.flush_interval_ms / 1000.0}s")
                     time.sleep(self.flush_interval_ms / 1000.0)
                 except Exception as e:
-                    logger.error(f"Error in auto-flush: {e}")
+                    logger.error(f"Error in auto-flush cycle #{cycle_count}: {e}")
+                    logger.debug(f"Auto-flush cycle #{cycle_count} error handled, sleeping {self.flush_interval_ms / 1000.0}s")
                     time.sleep(self.flush_interval_ms / 1000.0)
             
-            # Final flush on shutdown
-            if self.shutdown_requested:
-                try:
-                    loop.run_until_complete(self.flush())
-                except Exception as e:
-                    logger.error(f"Error in final flush: {e}")
-                finally:
+            logger.info(f"Auto-flush worker thread stopping (shutdown requested). Completed {cycle_count} cycles.")
+            
+            # Don't do final flush here - shutdown() will handle it with synchronous send
+            # This avoids event loop shutdown issues
+            logger.debug("Auto-flush thread skipping final flush (will be handled by shutdown() with sync send)")
+            
+            # Close the event loop
+            try:
+                if not loop.is_closed():
                     loop.close()
+                logger.debug("Auto-flush worker thread event loop closed")
+            except Exception:
+                pass  # Ignore errors during cleanup
 
-        flush_thread = threading.Thread(target=flush_worker, daemon=True)
+        flush_thread = threading.Thread(target=flush_worker, daemon=True, name="AIQA-AutoFlush")
         flush_thread.start()
         self.flush_timer = flush_thread
+        logger.info(f"Auto-flush thread started: {flush_thread.name} (daemon={flush_thread.daemon})")
 
     async def _send_spans(self, spans: List[Dict[str, Any]]) -> None:
-        """Send spans to the server API."""
-        if not self.server_url:
-            raise ValueError("AIQA_SERVER_URL is not set. Cannot send spans to server.")
-
+        """Send spans to the server API (async)."""
         import aiohttp
 
-        logger.debug(f"Sending {len(spans)} spans to server: {self.server_url}")
+        url = self._get_span_url()
+        headers = self._build_request_headers()
+        logger.debug(f"_send_spans() sending {len(spans)} spans to {url}")
+        if self.api_key:
+            logger.debug("_send_spans() using API key authentication")
+        else:
+            logger.debug("_send_spans() no API key provided")
 
-        headers = {
-            "Content-Type": "application/json",
-        }
+        try:
+            # Pre-serialize JSON to bytes and wrap in BytesIO to avoid blocking event loop
+            json_bytes = json.dumps(spans).encode('utf-8')
+            data = io.BytesIO(json_bytes)
+            
+            async with aiohttp.ClientSession() as session:
+                logger.debug(f"_send_spans() POST request starting to {url}")
+                async with session.post(url, data=data, headers=headers) as response:
+                    logger.debug(f"_send_spans() received response: status={response.status}")
+                    if not response.ok:
+                        error_text = await response.text()
+                        logger.error(
+                            f"_send_spans() failed: status={response.status}, "
+                            f"reason={response.reason}, error={error_text[:200]}"
+                        )
+                        raise Exception(
+                            f"Failed to send spans: {response.status} {response.reason} - {error_text}"
+                        )
+                    logger.debug(f"_send_spans() successfully sent {len(spans)} spans")
+        except RuntimeError as e:
+            if self._is_interpreter_shutdown_error(e):
+                if self.shutdown_requested:
+                    logger.debug(f"_send_spans() skipped due to interpreter shutdown: {e}")
+                else:
+                    logger.warning(f"_send_spans() interrupted by interpreter shutdown: {e}")
+                raise
+            logger.error(f"_send_spans() RuntimeError: {type(e).__name__}: {e}")
+            raise
+        except Exception as e:
+            logger.error(f"_send_spans() exception: {type(e).__name__}: {e}")
+            raise
+
+    def _send_spans_sync(self, spans: List[Dict[str, Any]]) -> None:
+        """Send spans to the server API (synchronous, for shutdown scenarios)."""
+        import requests
+
+        url = self._get_span_url()
+        headers = self._build_request_headers()
+        logger.debug(f"_send_spans_sync() sending {len(spans)} spans to {url}")
         if self.api_key:
-            headers["Authorization"] = f"ApiKey {self.api_key}"
+            logger.debug("_send_spans_sync() using API key authentication")
+        else:
+            logger.debug("_send_spans_sync() no API key provided")
 
-        async with aiohttp.ClientSession() as session:
-            async with session.post(
-                f"{self.server_url}/span",
-                json=spans,
-                headers=headers,
-            ) as response:
-                if not response.ok:
-                    error_text = await response.text()
-                    raise Exception(
-                        f"Failed to send spans: {response.status} {response.reason} - {error_text}"
-                    )
+        try:
+            response = requests.post(url, json=spans, headers=headers, timeout=10.0)
+            logger.debug(f"_send_spans_sync() received response: status={response.status_code}")
+            if not response.ok:
+                error_text = response.text[:200] if response.text else ""
+                logger.error(
+                    f"_send_spans_sync() failed: status={response.status_code}, "
+                    f"reason={response.reason}, error={error_text}"
+                )
+                raise Exception(
+                    f"Failed to send spans: {response.status_code} {response.reason} - {error_text}"
+                )
+            logger.debug(f"_send_spans_sync() successfully sent {len(spans)} spans")
+        except Exception as e:
+            logger.error(f"_send_spans_sync() exception: {type(e).__name__}: {e}")
+            raise
 
     def shutdown(self) -> None:
         """Shutdown the exporter, flushing any remaining spans. Call before process exit."""
+        logger.info("shutdown() called - initiating exporter shutdown")
         self.shutdown_requested = True
 
+        # Check buffer state before shutdown
+        with self.buffer_lock:
+            buffer_size = len(self.buffer)
+            logger.info(f"shutdown() buffer contains {buffer_size} span(s) before shutdown")
+
         # Wait for flush thread to finish (it will do final flush)
         if self.flush_timer and self.flush_timer.is_alive():
+            logger.info("shutdown() waiting for auto-flush thread to complete (timeout=10s)")
             self.flush_timer.join(timeout=10.0)
+            if self.flush_timer.is_alive():
+                logger.warning("shutdown() auto-flush thread did not complete within timeout")
+            else:
+                logger.info("shutdown() auto-flush thread completed")
+        else:
+            logger.debug("shutdown() no active auto-flush thread to wait for")
 
-        # Final flush attempt (synchronous)
-        import asyncio
-        try:
-            loop = asyncio.get_event_loop()
-            if loop.is_running():
-                # If loop is running, schedule flush
-                import concurrent.futures
-                with concurrent.futures.ThreadPoolExecutor() as executor:
-                    future = executor.submit(asyncio.run, self.flush())
-                    future.result(timeout=10.0)
+        # Final flush attempt (use synchronous send to avoid event loop issues)
+        with self.flush_lock:
+            logger.debug("shutdown() performing final flush with synchronous send")
+            # Atomically extract and remove spans to prevent race conditions
+            spans_to_flush = self._extract_and_remove_spans_from_buffer()
+            logger.debug(f"shutdown() extracted {len(spans_to_flush)} span(s) from buffer for final flush")
+
+            if spans_to_flush:
+                if not self.server_url:
+                    logger.warning(
+                        f"shutdown() skipping final flush: AIQA_SERVER_URL is not set. "
+                        f"{len(spans_to_flush)} span(s) will not be sent."
+                    )
+                    # Spans already removed from buffer, clear their keys to free memory
+                    self._remove_span_keys_from_tracking(spans_to_flush)
+                else:
+                    logger.info(f"shutdown() sending {len(spans_to_flush)} span(s) to server (synchronous)")
+                    try:
+                        self._send_spans_sync(spans_to_flush)
+                        logger.info(f"shutdown() successfully sent {len(spans_to_flush)} span(s) to server")
+                        # Spans already removed from buffer during extraction
+                        # Clear their keys from tracking set to free memory
+                        self._remove_span_keys_from_tracking(spans_to_flush)
+                    except Exception as e:
+                        logger.error(f"shutdown() failed to send spans: {e}")
+                        # Spans already removed, but process is exiting anyway
+                        logger.warning(f"shutdown() {len(spans_to_flush)} span(s) were not sent due to error")
+                        # Keys will remain in tracking set, but process is exiting so memory will be freed
             else:
-                loop.run_until_complete(self.flush())
-        except RuntimeError:
-            # No event loop, create one
-            asyncio.run(self.flush())
+                logger.debug("shutdown() no spans to flush")
+        
+        # Check buffer state after shutdown
+        with self.buffer_lock:
+            buffer_size = len(self.buffer)
+            if buffer_size > 0:
+                logger.warning(f"shutdown() buffer still contains {buffer_size} span(s) after shutdown")
+            else:
+                logger.info("shutdown() buffer is empty after shutdown")
+        
+        logger.info("shutdown() completed")