synth-ai 0.1.9__py3-none-any.whl → 0.2.1.dev0__py3-none-any.whl

synth_ai/tracing/events/manage.py

@@ -0,0 +1,147 @@
+import time
+from typing import Literal, Optional
+
+from synth_ai.tracing.abstractions import Event
+from synth_ai.tracing.events.store import event_store
+from synth_ai.tracing.local import _local, logger
+
+
+def get_current_event(event_type: str) -> "Event":
+    """
+    Get the current active event of the specified type.
+    Raises ValueError if no such event exists.
+    """
+    events = getattr(_local, "active_events", {})
+    if event_type not in events:
+        raise ValueError(f"No active event of type '{event_type}' found")
+    return events[event_type]
+
+
+def set_current_event(event: Optional["Event"], decorator_type: Literal["sync", "async"] = None):
+    """
+    Set the current event, ending any existing events of the same type.
+    If event is None, it clears the current event of that type.
+    """
+    if event is None:
+        raise ValueError("Event cannot be None when setting current event.")
+
+    # Check if we're in an async context
+    try:
+        import asyncio
+
+        asyncio.get_running_loop()
+        is_async = True
+    except RuntimeError:
+        is_async = False
+
+    if decorator_type == "sync" or not is_async:
+        # Original thread-local storage logic
+        if not hasattr(_local, "active_events"):
+            _local.active_events = {}
+            logger.debug("Initialized active_events in thread local storage")
+
+        # If there's an existing event of the same type, end it
+        if event.event_type in _local.active_events:
+            if (
+                _local.active_events[event.event_type].system_instance_id
+                == event.system_instance_id
+            ):
+                logger.debug(f"Found existing event of type {event.event_type}")
+                existing_event = _local.active_events[event.event_type]
+                existing_event.closed = time.time()
+                logger.debug(
+                    f"Closed existing event of type {event.event_type} at {existing_event.closed}"
+                )
+
+                # Store the closed event if system_instance_id is present
+                if hasattr(_local, "system_instance_id"):
+                    logger.debug(f"Storing closed event for system {_local.system_instance_id}")
+                    try:
+                        event_store.add_event(
+                            _local.system_name,
+                            _local.system_id,
+                            _local.system_instance_id,
+                            existing_event,
+                        )
+                        logger.debug("Successfully stored closed event")
+                    except Exception as e:
+                        logger.error(f"Failed to store closed event: {str(e)}")
+                        raise
+
+        # Set the new event with both keys
+        _local.active_events[event.event_type] = event  # Plain key
+        if is_async:
+            unique_key = f"{event.event_type}-{time.time()}"
+            _local.active_events[unique_key] = event  # Unique key for async
+
+    else:
+        from synth_ai.tracing.local import (
+            active_events_var,
+            system_id_var,
+            system_instance_id_var,
+            system_name_var,
+        )
+
+        # Get current active events from context var
+        active_events = active_events_var.get()
+        # print("Active events:", active_events)
+        # If there's an existing event of the same type, end it
+        if event and event.event_type in active_events:
+            existing_event = active_events[event.event_type]
+
+            # Check that the active event has the same system_instance_id as the one we're settting
+            if existing_event.system_instance_id == event.system_instance_id:
+                logger.debug(f"Found existing event of type {event.event_type}")
+                existing_event.closed = time.time()
+                logger.debug(
+                    f"Closed existing event of type {event.event_type} at {existing_event.closed}"
+                )
+
+                # Store the closed event if system_instance_id is present
+                system_instance_id = system_instance_id_var.get()
+                if system_instance_id:
+                    logger.debug(f"Storing closed event for system {system_instance_id}")
+                    try:
+                        event_store.add_event(
+                            system_name_var.get(),
+                            system_id_var.get(),
+                            system_instance_id,
+                            existing_event,
+                        )
+                        logger.debug("Successfully stored closed event")
+                    except Exception as e:
+                        logger.error(f"Failed to store closed event: {str(e)}")
+                        raise
+        elif not event:
+            logger.debug(f"No event found for type {event.event_type}")
+        # Set the new event with both keys
+        active_events[event.event_type] = event  # Plain key
+        unique_key = f"{event.event_type}-{time.time()}"
+        active_events[unique_key] = event  # Unique key for async
+        active_events_var.set(active_events)
+        logger.debug(
+            f"New event set as current in context vars with keys '{event.event_type}' and '{unique_key}'"
+        )
+
+
+def clear_current_event(event_type: str):
+    if hasattr(_local, "active_events"):
+        _local.active_events.pop(event_type, None)
+        logger.debug(f"Cleared current event of type {event_type}")
+
+
+def end_event(event_type: str) -> Optional[Event]:
+    """End the current event and store it."""
+    current_event = get_current_event(event_type)
+    if current_event:
+        current_event.closed = time.time()
+        # Store the event
+        if hasattr(_local, "system_instance_id"):
+            event_store.add_event(
+                _local.system_name,
+                _local.system_id,
+                _local.system_instance_id,
+                current_event,
+            )
+        clear_current_event(event_type)
+    return current_event

synth_ai/tracing/events/scope.py

@@ -0,0 +1,86 @@
+import time
+from contextlib import contextmanager
+
+from synth_ai.tracing.abstractions import Event
+from synth_ai.tracing.config import LoggingMode
+from synth_ai.tracing.decorators import (
+    _local,
+    clear_current_event,
+    get_tracing_config,
+    set_current_event,
+)
+from synth_ai.tracing.events.store import event_store
+from synth_ai.tracing.immediate_client import (
+    AsyncImmediateLogClient,
+    ImmediateLogClient,
+)
+from synth_ai.tracing.local import (
+    system_id_var,
+    system_instance_id_var,
+    system_name_var,
+)
+
+
+@contextmanager
+def event_scope(event_type: str):
+    """
+    Context manager for creating and managing events.
+
+    Usage:
+        with event_scope("my_event_type"):
+            # do stuff
+    """
+    # Check if we're in an async context
+    try:
+        import asyncio
+
+        asyncio.get_running_loop()
+        is_async = True
+    except RuntimeError:
+        is_async = False
+
+    # Get system_instance_id from appropriate source
+    system_instance_id = (
+        system_instance_id_var.get() if is_async else getattr(_local, "system_instance_id", None)
+    )
+    system_id = system_id_var.get() if is_async else getattr(_local, "system_id", None)
+    system_name = system_name_var.get() if is_async else getattr(_local, "system_name", None)
+
+    event = Event(
+        system_instance_id=system_instance_id,
+        event_type=event_type,
+        opened=time.time(),
+        closed=None,
+        partition_index=0,
+        agent_compute_step=None,
+        environment_compute_steps=[],
+    )
+    set_current_event(event)
+
+    try:
+        yield event
+    finally:
+        event.closed = time.time()
+        clear_current_event(event_type)
+
+        # Get the config to determine logging mode
+        config = get_tracing_config()
+
+        # If immediate logging is enabled and we have system info, send the event now
+        if config.mode == LoggingMode.INSTANT and system_instance_id:
+            system_info = {
+                "system_name": system_name,
+                "system_id": system_id,
+                "system_instance_id": system_instance_id,
+            }
+            if is_async:
+                client = AsyncImmediateLogClient(config)
+                # Note: Since we can't use await in a finally block,
+                # we'll have to rely on the event store as primary storage in async context
+            else:
+                client = ImmediateLogClient(config)
+                client.send_event(event, system_info)
+
+        # Always store in event_store
+        if system_instance_id:
+            event_store.add_event(system_name, system_id, system_instance_id, event)

synth_ai/tracing/events/store.py

@@ -0,0 +1,227 @@
+import json
+import logging
+import time
+from threading import RLock  # Change this import
+from typing import Any, Dict, List
+
+from synth_ai.tracing.abstractions import Event, EventPartitionElement, SystemTrace
+from synth_ai.tracing.local import (  # Import context variables
+    _local,
+    active_events_var,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class EventStore:
+    def __init__(self) -> None:
+        self._traces: Dict[str, SystemTrace] = {}
+        self._lock = RLock()  # Use RLock instead of Lock
+        self.logger = logging.getLogger(__name__)
+
+    def get_or_create_system_trace(
+        self,
+        system_name: str,
+        system_id: str,
+        system_instance_id: str,
+        _already_locked: bool = False,
+    ) -> SystemTrace:
+        """Get or create a SystemTrace for the given system_instance_id."""
+        logger = logging.getLogger(__name__)
+        # logger.debug(f"Starting get_or_create_system_trace for {system_instance_id}")
+
+        def _get_or_create():
+            # logger.debug("Inside _get_or_create")
+            if system_instance_id not in self._traces:
+                # Get system_instance_metadata from context if available
+                system_instance_metadata = {}
+                try:
+                    from synth_ai.tracing.context import get_current_context
+
+                    context = get_current_context()
+                    if (
+                        "system_instance_metadata" in context
+                        and context["system_instance_metadata"]
+                    ):
+                        system_instance_metadata = context["system_instance_metadata"]
+                except Exception as e:
+                    self.logger.warning(
+                        f"Failed to get system_instance_metadata from context: {str(e)}"
+                    )
+
+                # logger.debug(f"Creating new system trace for {system_instance_id}")
+                self._traces[system_instance_id] = SystemTrace(
+                    system_name=system_name,
+                    system_id=system_id,
+                    system_instance_id=system_instance_id,
+                    metadata={},
+                    partition=[],  # EventPartitionElement(partition_index=0, events=[])
+                    current_partition_index=0,
+                    instance_metadata=system_instance_metadata,
+                )
+            # logger.debug("Returning system trace")
+            return self._traces[system_instance_id]
+
+        if _already_locked:
+            return _get_or_create()
+        else:
+            with self._lock:
+                # logger.debug("Lock acquired in get_or_create_system_trace")
+                return _get_or_create()
+
+    def increment_partition(self, system_name: str, system_id: str, system_instance_id: str) -> int:
+        """Increment the partition index for a system and create new partition element."""
+        logger = logging.getLogger(__name__)
+        # logger.debug(f"Starting increment_partition for system {system_instance_id}")
+
+        with self._lock:
+            # logger.debug("Lock acquired in increment_partition")
+            system_trace = self.get_or_create_system_trace(
+                system_name, system_id, system_instance_id, _already_locked=True
+            )
+            # logger.debug(
+            #     f"Got system trace, current index: {system_trace.current_partition_index}"
+            # )
+
+            system_trace.current_partition_index += 1
+            # logger.debug(
+            #     f"Incremented index to: {system_trace.current_partition_index}"
+            # )
+
+            system_trace.partition.append(
+                EventPartitionElement(
+                    partition_index=system_trace.current_partition_index, events=[]
+                )
+            )
+            # logger.debug("Added new partition element")
+
+            return system_trace.current_partition_index
+
+    def add_event(self, system_name: str, system_id: str, system_instance_id: str, event: Event):
+        """Add an event to the appropriate partition of the system trace."""
+        # self.#logger.debug(f"Adding event type {event.event_type} to system {system_instance_id}")
+        # self.#logger.debug(
+        #     f"Event details: opened={event.opened}, closed={event.closed}, partition={event.partition_index}"
+        # )
+        # print("Adding event: ", event)
+
+        # try:
+        if not self._lock.acquire(timeout=5):
+            self.logger.error("Failed to acquire lock within timeout period")
+            return
+
+        try:
+            system_trace = self.get_or_create_system_trace(
+                system_name, system_id, system_instance_id, _already_locked=True
+            )
+            # self.#logger.debug(
+            #     f"Got system trace with {len(system_trace.partition)} partitions"
+            # )
+
+            current_partition = next(
+                (p for p in system_trace.partition if p.partition_index == event.partition_index),
+                None,
+            )
+
+            if current_partition is None:
+                self.logger.error(
+                    f"No partition found for index {event.partition_index} - existing partitions: {set([p.partition_index for p in system_trace.partition])}"
+                )
+                raise ValueError(f"No partition found for index {event.partition_index}")
+
+            current_partition.events.append(event)
+        finally:
+            self._lock.release()
+        # except Exception as e:
+        #     self.logger.error(f"Error in add_event: {str(e)}", exc_info=True)
+        #     raise
+
+    def get_system_traces(self) -> List[SystemTrace]:
+        """Get all system traces."""
+        with self._lock:
+            self.end_all_active_events()
+
+            return list(self._traces.values())
+
+    def end_all_active_events(self):
+        """End all active events and store them."""
+        # self.#logger.debug("Ending all active events")
+
+        # For synchronous code
+        if hasattr(_local, "active_events"):
+            for _, event in list(_local.active_events.items()):
+                if event.closed is None:
+                    event.closed = time.time()
+                    self.add_event(
+                        event.system_name,
+                        event.system_id,
+                        event.system_instance_id,
+                        event,
+                    )
+            _local.active_events.clear()
+
+        else:
+            # For asynchronous code
+            active_events_async = active_events_var.get()
+            if active_events_async:
+                for _, event in list(active_events_async.items()):
+                    if event.closed is None:
+                        event.closed = time.time()
+                        self.add_event(
+                            event.system_name,
+                            event.system_id,
+                            event.system_instance_id,
+                            event,
+                        )
+                active_events_var.set({})
+
+    def get_system_traces_json(self) -> str:
+        """Get all system traces as JSON."""
+        with self._lock:
+            return json.dumps(
+                [
+                    {
+                        "system_instance_id": trace.system_instance_id,
+                        "current_partition_index": trace.current_partition_index,
+                        "partition": [
+                            {
+                                "partition_index": p.partition_index,
+                                "events": [self._event_to_dict(event) for event in p.events],
+                            }
+                            for p in trace.partition
+                        ],
+                    }
+                    for trace in self._traces.values()
+                ],
+                default=str,
+            )
+
+    def _event_to_dict(self, event: Event) -> Dict[str, Any]:
+        """Convert an Event object to a dictionary."""
+        return {
+            "event_type": event.event_type,
+            "opened": event.opened,
+            "closed": event.closed,
+            "partition_index": event.partition_index,
+            "agent_compute_step": {
+                "event_order": event.agent_compute_step.event_order,
+                "compute_began": event.agent_compute_step.compute_began,
+                "compute_ended": event.agent_compute_step.compute_ended,
+                "compute_input": event.agent_compute_step.compute_input,
+                "compute_output": event.agent_compute_step.compute_output,
+            },
+            "environment_compute_steps": [
+                {
+                    "event_order": step.event_order,
+                    "compute_began": step.compute_began,
+                    "compute_ended": step.compute_ended,
+                    "compute_input": step.compute_input,
+                    "compute_output": step.compute_output,
+                }
+                for step in event.environment_compute_steps
+            ],
+        }
+
+
+# Global event store instance
+event_store = EventStore()

synth_ai/tracing/immediate_client.py

@@ -0,0 +1,152 @@
+import asyncio
+import logging
+import time
+from typing import Dict
+
+import httpx
+
+from synth_ai.tracing.abstractions import Event
+from synth_ai.tracing.client_manager import ClientManager
+from synth_ai.tracing.config import TracingConfig
+from synth_ai.tracing.log_client_base import BaseAsyncLogClient, BaseLogClient
+
+logger = logging.getLogger(__name__)
+
+
+class ImmediateLogClient(BaseLogClient):
+    """Synchronous client for immediate logging of events"""
+
+    def __init__(self, config: TracingConfig):
+        super().__init__(config)
+        self.client_manager = ClientManager.initialize(config)
+
+    def send_event(self, event: Event, system_info: Dict[str, str]) -> bool:
+        """Send a single event with retries and fallback"""
+        from synth_ai.tracing.retry_queue import (
+            retry_queue,  # Import here to avoid circular import
+        )
+
+        payload = self._prepare_payload(event, system_info)
+        headers = {"Authorization": f"Bearer {self.config.api_key}"}
+        last_exception = None
+
+        for attempt in range(self.config.max_retries):
+            try:
+                response = self.client_manager.sync_client.post(
+                    f"{self.config.base_url}/v1/uploads/stream",
+                    json=payload,
+                    headers=headers,
+                    timeout=self.config.timeout,
+                )
+                logger.info(f"Event upload response status: {response.status_code}")
+                response.raise_for_status()
+                response_data = response.json()
+                event.id = response_data.get("event_id")  # Store the event_id
+                self._handle_success()
+                return True
+            except Exception as e:
+                last_exception = e
+                status_code = getattr(e, "response", None)
+                if status_code:
+                    status_code = status_code.status_code
+
+                if not self._should_retry(attempt, status_code):
+                    break
+
+                backoff = self.client_manager.calculate_backoff(attempt)
+                time.sleep(backoff)
+
+        # If we get here, all immediate retries failed
+        self._handle_failure(event, system_info, last_exception)
+
+        # Add to retry queue for later processing
+        retry_queue.add_failed_event(event, system_info)
+        return False
+
+
+class AsyncImmediateLogClient(BaseAsyncLogClient):
+    """Asynchronous client for immediate logging of events"""
+
+    def __init__(self, config: TracingConfig):
+        super().__init__(config)
+        self.client_manager = ClientManager.initialize(config)
+
+    async def send_event(self, event: Event, system_info: Dict[str, str]) -> bool:
+        """Send a single event with retries and fallback (async version)"""
+        from synth_ai.tracing.retry_queue import retry_queue
+
+        if not self.config.api_key:
+            logger.error("No API key provided")
+            return False
+
+        # First get JWT token
+        async with httpx.AsyncClient(timeout=self.config.timeout) as client:
+            try:
+                auth_response = await client.get(
+                    f"{self.config.base_url}/v1/auth/token",
+                    headers={"customer_specific_api_key": self.config.api_key.strip()},
+                )
+                auth_response.raise_for_status()
+                auth_data = auth_response.json()  # This is synchronous in httpx
+                logger.debug(f"Auth response data: {auth_data}")
+
+                token = auth_data.get("access_token")
+                if not token:
+                    logger.error(
+                        f"No access token received from auth endpoint. Response data: {auth_data}"
+                    )
+                    return False
+            except Exception as e:
+                # logger.error(f"Failed to get auth token: {e}")
+                return False
+
+            # Now send the event with the JWT token
+            payload = self._prepare_payload(event, system_info)
+            headers = {
+                "Authorization": f"Bearer {token}",
+                "Content-Type": "application/json",
+                "Accept": "application/json",
+            }
+
+            logger.debug(f"Request URL: {self.config.base_url}/v1/uploads/stream")
+            logger.debug("Using JWT token for authentication")
+            logger.debug(f"Headers: {headers}")
+            logger.debug(f"Payload size: {len(str(payload))} bytes")
+
+            async with httpx.AsyncClient(timeout=self.config.timeout) as client:
+                for attempt in range(self.config.max_retries + 1):
+                    try:
+                        response = await client.post(
+                            f"{self.config.base_url}/v1/uploads/stream",
+                            json=payload,
+                            headers=headers,
+                            timeout=self.config.timeout,
+                        )
+                        logger.info(f"Event upload response status: {response.status_code}")
+
+                        if response.status_code >= 500:
+                            error_body = response.text
+                            logger.error(f"Server error. Response: {error_body}")
+                            response.raise_for_status()
+
+                        if response.status_code == 401:
+                            error_body = response.text
+                            logger.error(f"Authentication failed. Response: {error_body}")
+                            response.raise_for_status()
+
+                        response_data = response.json()  # This is synchronous in httpx
+                        event.id = response_data.get("event_id")
+                        return True
+
+                    except Exception as e:
+                        # last_exception = e
+                        # logger.error(f"Upload attempt {attempt + 1} failed: {str(e)}")
+                        if attempt < self.config.max_retries:
+                            backoff = self.client_manager.calculate_backoff(attempt)
+                            await asyncio.sleep(backoff)
+
+                # logger.error(
+                #     f"All upload attempts failed. Last error: {last_exception}"
+                # )
+                retry_queue.add_failed_event(event, system_info)
+                return False

synth_ai/tracing/local.py

@@ -0,0 +1,18 @@
+import logging
+import threading
+from contextvars import ContextVar
+from typing import Any, Dict
+
+logger = logging.getLogger(__name__)
+
+# Thread-local storage for active events and system_instance_id
+# Used for synchronous tracing
+_local = threading.local()
+# Used for asynchronous tracing
+system_name_var: ContextVar[str] = ContextVar("system_name", default=None)
+system_id_var: ContextVar[str] = ContextVar("system_id", default=None)
+system_instance_id_var: ContextVar[str] = ContextVar("system_instance_id", default=None)
+system_instance_metadata_var: ContextVar[Dict[str, Any]] = ContextVar(
+    "system_instance_metadata", default={}
+)
+active_events_var: ContextVar[Dict[str, Any]] = ContextVar("active_events", default={})