fluxloop 0.1.0__py3-none-any.whl

fluxloop/context.py

@@ -0,0 +1,219 @@
+"""
+Context management for tracing.
+"""
+
+import contextvars
+import random
+from contextlib import contextmanager
+from datetime import datetime, timezone
+from typing import Any, Dict, List, Optional, Union
+from uuid import UUID, uuid4
+
+from .buffer import EventBuffer
+from .config import get_config
+from .models import ObservationData, TraceData
+
+
+# Context variable for current FluxLoop context
+_context_var: contextvars.ContextVar[Optional["FluxLoopContext"]] = contextvars.ContextVar(
+    "fluxloop_context", default=None
+)
+
+
+class FluxLoopContext:
+    """
+    Context for managing trace and observation hierarchy.
+    """
+    
+    def __init__(
+        self,
+        trace_name: str,
+        trace_id: Optional[UUID] = None,
+        session_id: Optional[UUID] = None,
+        user_id: Optional[str] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+        tags: Optional[List[str]] = None,
+        trace_id_override: Optional[Union[UUID, str]] = None,
+    ):
+        """
+        Initialize a new context.
+        
+        Args:
+            trace_name: Name for the trace
+            trace_id: Optional trace ID (generated if not provided)
+            session_id: Optional session ID for grouping traces
+            user_id: Optional user identifier
+            metadata: Additional metadata
+            tags: Tags for categorization
+        """
+        self.config = get_config()
+        self.buffer = EventBuffer.get_instance()
+        
+        # Create trace
+        trace_uuid = self._coerce_uuid(trace_id_override) or trace_id or uuid4()
+
+        self.trace = TraceData(
+            id=trace_uuid,
+            name=trace_name,
+            session_id=session_id,
+            user_id=user_id,
+            metadata=metadata or {},
+            tags=tags or [],
+        )
+        
+        # Observation stack for hierarchy
+        self.observation_stack: List[ObservationData] = []
+        self.observations: List[ObservationData] = []
+        
+        # Sampling decision
+        self.is_sampled = random.random() < self.config.sample_rate
+    
+    def is_enabled(self) -> bool:
+        """Check if tracing is enabled and sampled."""
+        return self.config.enabled and self.is_sampled
+    
+    def push_observation(self, observation: ObservationData) -> None:
+        """
+        Push a new observation onto the stack.
+        
+        Args:
+            observation: Observation to push
+        """
+        if not self.is_enabled():
+            return
+        
+        # Set trace ID
+        observation.trace_id = self.trace.id
+        
+        # Set parent if there's an observation on the stack
+        if self.observation_stack:
+            parent = self.observation_stack[-1]
+            observation.parent_observation_id = parent.id
+        
+        # Add to stack and list
+        self.observation_stack.append(observation)
+        self.observations.append(observation)
+        
+        # Send to buffer if it's complete (has end_time)
+        if observation.end_time:
+            self.buffer.add_observation(self.trace.id, observation)
+    
+    def pop_observation(self) -> Optional[ObservationData]:
+        """
+        Pop the current observation from the stack.
+        
+        Returns:
+            The popped observation, or None if stack is empty
+        """
+        if not self.is_enabled() or not self.observation_stack:
+            return None
+        
+        observation = self.observation_stack.pop()
+        
+        # If observation doesn't have end_time, set it now
+        if not observation.end_time:
+            observation.end_time = datetime.now(timezone.utc)
+        
+        # Send to buffer
+        self.buffer.add_observation(self.trace.id, observation)
+        
+        return observation
+    
+    def add_metadata(self, key: str, value: Any) -> None:
+        """Add metadata to the current trace."""
+        if self.is_enabled():
+            self.trace.metadata[key] = value
+    
+    def add_tag(self, tag: str) -> None:
+        """Add a tag to the current trace."""
+        if self.is_enabled() and tag not in self.trace.tags:
+            self.trace.tags.append(tag)
+    
+    @staticmethod
+    def _coerce_uuid(value: Optional[Union[UUID, str]]) -> Optional[UUID]:
+        if value is None:
+            return None
+        if isinstance(value, UUID):
+            return value
+        try:
+            return UUID(str(value))
+        except Exception:
+            return None
+
+    def set_user(self, user_id: str) -> None:
+        """Set the user ID for the trace."""
+        if self.is_enabled():
+            self.trace.user_id = user_id
+    
+    def finalize(self) -> None:
+        """Finalize the trace and send all data."""
+        if not self.is_enabled():
+            return
+        
+        # Pop any remaining observations
+        while self.observation_stack:
+            self.pop_observation()
+        
+        # Set trace end time
+        if not self.trace.end_time:
+            self.trace.end_time = datetime.now(timezone.utc)
+        
+        # Send trace to buffer
+        self.buffer.add_trace(self.trace)
+        
+        # Trigger flush if needed
+        self.buffer.flush_if_needed()
+
+
+def get_current_context() -> Optional[FluxLoopContext]:
+    """
+    Get the current FluxLoop context.
+    
+    Returns:
+        Current context or None if not in a context
+    """
+    return _context_var.get()
+
+
+@contextmanager
+def instrument(
+    name: str,
+    session_id: Optional[UUID] = None,
+    user_id: Optional[str] = None,
+    metadata: Optional[Dict[str, Any]] = None,
+    tags: Optional[List[str]] = None,
+    trace_id: Optional[UUID] = None,
+):
+    """
+    Context manager for instrumenting code blocks.
+    
+    Args:
+        name: Name for the trace
+        session_id: Optional session ID
+        user_id: Optional user identifier
+        metadata: Additional metadata
+        tags: Tags for categorization
+        
+    Example:
+        >>> with fluxloop.instrument("my_workflow"):
+        ...     result = my_agent.process(input_data)
+    """
+    # Create new context
+    context = FluxLoopContext(
+        trace_name=name,
+        session_id=session_id,
+        user_id=user_id,
+        metadata=metadata,
+        tags=tags,
+        trace_id_override=trace_id,
+    )
+    
+    # Set as current context
+    token = _context_var.set(context)
+    
+    try:
+        yield context
+    finally:
+        # Finalize and reset context
+        context.finalize()
+        _context_var.reset(token)

fluxloop/decorators.py

@@ -0,0 +1,465 @@
+"""
+Decorators for instrumenting agent code.
+"""
+
+import functools
+import inspect
+import time
+import traceback
+from datetime import datetime, timezone
+from typing import Any, Callable, Dict, Optional, TypeVar, Union
+from uuid import UUID, uuid4
+
+from .context import get_current_context
+from .models import ObservationData, ObservationType
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+
+def agent(
+    name: Optional[str] = None,
+    metadata: Optional[Dict[str, Any]] = None,
+    capture_input: bool = True,
+    capture_output: bool = True,
+) -> Callable[[F], F]:
+    """
+    Decorator for agent entry points.
+    
+    Args:
+        name: Name for the agent trace (defaults to function name)
+        metadata: Additional metadata to attach
+        capture_input: Whether to capture function arguments
+        capture_output: Whether to capture return value
+        
+    Example:
+        >>> @fluxloop.agent(name="ChatBot")
+        ... def process_message(message: str) -> str:
+        ...     return f"Response to: {message}"
+    """
+    
+    def decorator(func: F) -> F:
+        agent_name = name or func.__name__
+        
+        @functools.wraps(func)
+        def sync_wrapper(*args, **kwargs):
+            context = get_current_context()
+            if not context or not context.is_enabled():
+                return func(*args, **kwargs)
+            
+            # Create observation
+            obs_id = uuid4()
+            start_time = datetime.now(timezone.utc)
+            
+            # Capture input
+            input_data = None
+            if capture_input:
+                input_data = _serialize_arguments(func, args, kwargs)
+            
+            # Create observation data
+            observation = ObservationData(
+                id=obs_id,
+                type=ObservationType.AGENT,
+                name=agent_name,
+                start_time=start_time,
+                input=input_data,
+                metadata=metadata or {},
+            )
+            
+            # Push to context
+            context.push_observation(observation)
+            
+            try:
+                # Execute function
+                result = func(*args, **kwargs)
+                
+                # Capture output
+                if capture_output:
+                    observation.output = _serialize_value(result)
+                
+                return result
+                
+            except Exception as e:
+                # Capture error
+                observation.error = str(e)
+                observation.metadata["error_type"] = type(e).__name__
+                observation.metadata["traceback"] = traceback.format_exc()
+                raise
+                
+            finally:
+                # Finalize observation
+                observation.end_time = datetime.now(timezone.utc)
+                context.pop_observation()
+        
+        @functools.wraps(func)
+        async def async_wrapper(*args, **kwargs):
+            context = get_current_context()
+            if not context or not context.is_enabled():
+                return await func(*args, **kwargs)
+            
+            # Create observation
+            obs_id = uuid4()
+            start_time = datetime.now(timezone.utc)
+            
+            # Capture input
+            input_data = None
+            if capture_input:
+                input_data = _serialize_arguments(func, args, kwargs)
+            
+            # Create observation data
+            observation = ObservationData(
+                id=obs_id,
+                type=ObservationType.AGENT,
+                name=agent_name,
+                start_time=start_time,
+                input=input_data,
+                metadata=metadata or {},
+            )
+            
+            # Push to context
+            context.push_observation(observation)
+            
+            try:
+                # Execute function
+                result = await func(*args, **kwargs)
+                
+                # Capture output
+                if capture_output:
+                    observation.output = _serialize_value(result)
+                
+                return result
+                
+            except Exception as e:
+                # Capture error
+                observation.error = str(e)
+                observation.metadata["error_type"] = type(e).__name__
+                observation.metadata["traceback"] = traceback.format_exc()
+                raise
+                
+            finally:
+                # Finalize observation
+                observation.end_time = datetime.now(timezone.utc)
+                context.pop_observation()
+        
+        # Return appropriate wrapper based on function type
+        if inspect.iscoroutinefunction(func):
+            return async_wrapper
+        else:
+            return sync_wrapper
+    
+    return decorator
+
+
+def prompt(
+    name: Optional[str] = None,
+    model: Optional[str] = None,
+    capture_tokens: bool = True,
+) -> Callable[[F], F]:
+    """
+    Decorator for prompt/LLM generation functions.
+    
+    Args:
+        name: Name for the generation (defaults to function name)
+        model: Model name being used
+        capture_tokens: Whether to try to capture token usage
+        
+    Example:
+        >>> @fluxloop.prompt(model="gpt-3.5-turbo")
+        ... def generate_response(prompt: str) -> str:
+        ...     return llm.generate(prompt)
+    """
+    
+    def decorator(func: F) -> F:
+        prompt_name = name or func.__name__
+        
+        @functools.wraps(func)
+        def sync_wrapper(*args, **kwargs):
+            context = get_current_context()
+            if not context or not context.is_enabled():
+                return func(*args, **kwargs)
+            
+            # Create observation
+            obs_id = uuid4()
+            start_time = datetime.now(timezone.utc)
+            
+            # Capture input
+            input_data = _serialize_arguments(func, args, kwargs)
+            
+            # Create observation data
+            observation = ObservationData(
+                id=obs_id,
+                type=ObservationType.GENERATION,
+                name=prompt_name,
+                start_time=start_time,
+                input=input_data,
+                model=model,
+                metadata={},
+            )
+            
+            # Push to context
+            context.push_observation(observation)
+            
+            try:
+                # Execute function
+                result = func(*args, **kwargs)
+                
+                # Capture output
+                observation.output = _serialize_value(result)
+                
+                # Try to extract token usage if result is a dict-like object
+                if capture_tokens and hasattr(result, "get"):
+                    if "usage" in result:
+                        usage = result["usage"]
+                        observation.prompt_tokens = usage.get("prompt_tokens")
+                        observation.completion_tokens = usage.get("completion_tokens")
+                        observation.total_tokens = usage.get("total_tokens")
+                
+                return result
+                
+            except Exception as e:
+                # Capture error
+                observation.error = str(e)
+                observation.metadata["error_type"] = type(e).__name__
+                observation.metadata["traceback"] = traceback.format_exc()
+                raise
+                
+            finally:
+                # Finalize observation
+                observation.end_time = datetime.now(timezone.utc)
+                context.pop_observation()
+        
+        @functools.wraps(func)
+        async def async_wrapper(*args, **kwargs):
+            context = get_current_context()
+            if not context or not context.is_enabled():
+                return await func(*args, **kwargs)
+            
+            # Create observation
+            obs_id = uuid4()
+            start_time = datetime.now(timezone.utc)
+            
+            # Capture input
+            input_data = _serialize_arguments(func, args, kwargs)
+            
+            # Create observation data
+            observation = ObservationData(
+                id=obs_id,
+                type=ObservationType.GENERATION,
+                name=prompt_name,
+                start_time=start_time,
+                input=input_data,
+                model=model,
+                metadata={},
+            )
+            
+            # Push to context
+            context.push_observation(observation)
+            
+            try:
+                # Execute function
+                result = await func(*args, **kwargs)
+                
+                # Capture output
+                observation.output = _serialize_value(result)
+                
+                # Try to extract token usage if result is a dict-like object
+                if capture_tokens and hasattr(result, "get"):
+                    if "usage" in result:
+                        usage = result["usage"]
+                        observation.prompt_tokens = usage.get("prompt_tokens")
+                        observation.completion_tokens = usage.get("completion_tokens")
+                        observation.total_tokens = usage.get("total_tokens")
+                
+                return result
+                
+            except Exception as e:
+                # Capture error
+                observation.error = str(e)
+                observation.metadata["error_type"] = type(e).__name__
+                observation.metadata["traceback"] = traceback.format_exc()
+                raise
+                
+            finally:
+                # Finalize observation
+                observation.end_time = datetime.now(timezone.utc)
+                context.pop_observation()
+        
+        # Return appropriate wrapper based on function type
+        if inspect.iscoroutinefunction(func):
+            return async_wrapper
+        else:
+            return sync_wrapper
+    
+    return decorator
+
+
+def tool(
+    name: Optional[str] = None,
+    description: Optional[str] = None,
+) -> Callable[[F], F]:
+    """
+    Decorator for tool/function calls.
+    
+    Args:
+        name: Name for the tool (defaults to function name)
+        description: Description of what the tool does
+        
+    Example:
+        >>> @fluxloop.tool(description="Search the web")
+        ... def web_search(query: str) -> List[str]:
+        ...     return search_engine.search(query)
+    """
+    
+    def decorator(func: F) -> F:
+        tool_name = name or func.__name__
+        
+        @functools.wraps(func)
+        def sync_wrapper(*args, **kwargs):
+            context = get_current_context()
+            if not context or not context.is_enabled():
+                return func(*args, **kwargs)
+            
+            # Create observation
+            obs_id = uuid4()
+            start_time = datetime.now(timezone.utc)
+            
+            # Capture input
+            input_data = _serialize_arguments(func, args, kwargs)
+            
+            # Create observation data
+            observation = ObservationData(
+                id=obs_id,
+                type=ObservationType.TOOL,
+                name=tool_name,
+                start_time=start_time,
+                input=input_data,
+                metadata={"description": description} if description else {},
+            )
+            
+            # Push to context
+            context.push_observation(observation)
+            
+            try:
+                # Execute function
+                result = func(*args, **kwargs)
+                
+                # Capture output
+                observation.output = _serialize_value(result)
+                
+                return result
+                
+            except Exception as e:
+                # Capture error
+                observation.error = str(e)
+                observation.metadata["error_type"] = type(e).__name__
+                observation.metadata["traceback"] = traceback.format_exc()
+                raise
+                
+            finally:
+                # Finalize observation
+                observation.end_time = datetime.now(timezone.utc)
+                context.pop_observation()
+        
+        @functools.wraps(func)
+        async def async_wrapper(*args, **kwargs):
+            context = get_current_context()
+            if not context or not context.is_enabled():
+                return await func(*args, **kwargs)
+            
+            # Create observation
+            obs_id = uuid4()
+            start_time = datetime.now(timezone.utc)
+            
+            # Capture input
+            input_data = _serialize_arguments(func, args, kwargs)
+            
+            # Create observation data
+            observation = ObservationData(
+                id=obs_id,
+                type=ObservationType.TOOL,
+                name=tool_name,
+                start_time=start_time,
+                input=input_data,
+                metadata={"description": description} if description else {},
+            )
+            
+            # Push to context
+            context.push_observation(observation)
+            
+            try:
+                # Execute function
+                result = await func(*args, **kwargs)
+                
+                # Capture output
+                observation.output = _serialize_value(result)
+                
+                return result
+                
+            except Exception as e:
+                # Capture error
+                observation.error = str(e)
+                observation.metadata["error_type"] = type(e).__name__
+                observation.metadata["traceback"] = traceback.format_exc()
+                raise
+                
+            finally:
+                # Finalize observation
+                observation.end_time = datetime.now(timezone.utc)
+                context.pop_observation()
+        
+        # Return appropriate wrapper based on function type
+        if inspect.iscoroutinefunction(func):
+            return async_wrapper
+        else:
+            return sync_wrapper
+    
+    return decorator
+
+
+def _serialize_arguments(func: Callable, args: tuple, kwargs: dict) -> Dict[str, Any]:
+    """Serialize function arguments for storage."""
+    sig = inspect.signature(func)
+    bound = sig.bind(*args, **kwargs)
+    bound.apply_defaults()
+    
+    serialized = {}
+    for param_name, param_value in bound.arguments.items():
+        serialized[param_name] = _serialize_value(param_value)
+    
+    return serialized
+
+
+def _serialize_value(value: Any) -> Any:
+    """Serialize a value for storage."""
+    # Handle common types
+    if value is None or isinstance(value, (str, int, float, bool)):
+        return value
+    
+    # Handle UUIDs
+    if isinstance(value, UUID):
+        return str(value)
+    
+    # Handle datetime
+    if isinstance(value, datetime):
+        return value.isoformat()
+    
+    # Handle lists/tuples
+    if isinstance(value, (list, tuple)):
+        return [_serialize_value(v) for v in value]
+    
+    # Handle dicts
+    if isinstance(value, dict):
+        return {k: _serialize_value(v) for k, v in value.items()}
+    
+    # Handle objects with dict representation
+    if hasattr(value, "dict"):
+        return value.dict()
+    
+    # Handle objects with model_dump (Pydantic v2)
+    if hasattr(value, "model_dump"):
+        return value.model_dump()
+    
+    # Fallback to string representation
+    try:
+        return str(value)
+    except:
+        return f"<{type(value).__name__}>"