epi-recorder 2.1.3__py3-none-any.whl → 2.3.0__py3-none-any.whl

epi_recorder/async_api.py

@@ -0,0 +1,151 @@
+import asyncio
+import threading
+import time
+from concurrent.futures import ThreadPoolExecutor
+from contextlib import asynccontextmanager
+from typing import Optional, Dict, Any
+from datetime import datetime
+
+from epi_core.storage import EpiStorage
+from epi_core.schemas import StepModel
+
+class AsyncRecorder:
+    """
+    Async-native recorder that doesn't block the event loop.
+    Uses background thread for SQLite writes.
+    """
+    
+    def __init__(self, session_name: str, output_dir: str = "."):
+        self.session_name = session_name
+        self.output_dir = output_dir
+        
+        # Thread-safe queue for steps
+        self._queue = asyncio.Queue()
+        
+        # Background thread executor (1 thread is enough for SQLite)
+        self._executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="epi_writer")
+        
+        # Storage instance (created in background thread)
+        self._storage: Optional[EpiStorage] = None
+        self._writer_task: Optional[asyncio.Task] = None
+        
+        # State tracking
+        self._step_count = 0
+        self._done = asyncio.Event()
+        self._error: Optional[Exception] = None
+    
+    async def start(self):
+        """Initialize storage in background thread and start writer"""
+        # Create storage in thread (SQLite init is also blocking)
+        loop = asyncio.get_event_loop()
+        self._storage = await loop.run_in_executor(
+            self._executor,
+            lambda: EpiStorage(self.session_name, self.output_dir)
+        )
+        
+        # Start background writer task
+        self._writer_task = asyncio.create_task(self._writer_loop())
+    
+    async def record_step(self, step_type: str, content: dict):
+        """Non-blocking step recording"""
+        if self._error:
+            raise self._error
+        
+        self._step_count += 1
+        
+        # Put in queue (never blocks, just buffers in memory)
+        await self._queue.put({
+            'index': self._step_count,
+            'type': step_type,
+            'content': content,
+            'timestamp': datetime.utcnow() # StepModel expects datetime
+        })
+    
+    async def _writer_loop(self):
+        """Background task: Drains queue to SQLite"""
+        try:
+            while True:
+                # Wait for item with timeout to check for shutdown
+                try:
+                    step_data = await asyncio.wait_for(self._queue.get(), timeout=0.5)
+                except asyncio.TimeoutError:
+                    continue
+                
+                if step_data is None: # Shutdown sentinel
+                   self._queue.task_done()
+                   break
+
+                # Write to SQLite in background thread (non-blocking for async)
+                loop = asyncio.get_event_loop()
+                await loop.run_in_executor(
+                    self._executor,
+                    self._write_to_storage,
+                    step_data
+                )
+                
+                self._queue.task_done()
+                
+        except asyncio.CancelledError:
+            # Graceful shutdown
+            pass
+        except Exception as e:
+            self._error = e
+    
+    def _write_to_storage(self, step_data: dict):
+        """Synchronous SQLite write (runs in background thread)"""
+        if self._storage:
+            # Construct StepModel
+            step = StepModel(
+                index=step_data['index'],
+                timestamp=step_data['timestamp'],
+                kind=step_data['type'],
+                content=step_data['content']
+            )
+            self._storage.add_step(step)
+    
+    async def stop(self):
+        """Finalize: Drain queue, close storage"""
+        if not self._writer_task:
+            return
+        
+        # Signal writer to finish
+        await self._queue.put(None)
+        await self._queue.join()
+        
+        # Wait for task
+        self._writer_task.cancel()
+        try:
+            await self._writer_task
+        except asyncio.CancelledError:
+            pass
+        
+        # Finalize storage in background thread
+        if self._storage:
+            loop = asyncio.get_event_loop()
+            await loop.run_in_executor(
+                self._executor,
+                self._storage.finalize
+            )
+        
+        # Shutdown executor
+        self._executor.shutdown(wait=True)
+
+@asynccontextmanager
+async def record_async(session_name: str, output_dir: str = "."):
+    """
+    Async context manager for recording.
+    
+    Usage:
+        async with record_async("my_agent") as rec:
+            await agent.arun("task")  # Non-blocking
+    """
+    recorder = AsyncRecorder(session_name, output_dir)
+    await recorder.start()
+    try:
+        yield recorder
+    finally:
+        await recorder.stop()
+
+
+
+ 

epi_recorder/bootstrap.py

@@ -56,3 +56,7 @@ def initialize_recording():
 # Auto-initialize if EPI_RECORD is set
 if os.environ.get("EPI_RECORD") == "1":
     initialize_recording()
+
+
+
+ 

epi_recorder/environment.py

@@ -235,3 +235,7 @@ def capture_environment(
         include_all_env_vars=include_all_env_vars,
         redact_env_vars=redact_env_vars
     )
+
+
+
+ 

epi_recorder/patcher.py

@@ -14,6 +14,7 @@ from functools import wraps
 
 from epi_core.schemas import StepModel
 from epi_core.redactor import get_default_redactor
+from epi_core.storage import EpiStorage
 
 
 class RecordingContext:
@@ -32,7 +33,6 @@ class RecordingContext:
             enable_redaction: Whether to redact secrets (default: True)
         """
         self.output_dir = output_dir
-        # self.steps: List[StepModel] = []  # Removed for scalability
         self.step_index = 0
         self.enable_redaction = enable_redaction
         self.redactor = get_default_redactor() if enable_redaction else None
@@ -40,9 +40,16 @@ class RecordingContext:
         # Ensure output directory exists
         self.output_dir.mkdir(parents=True, exist_ok=True)
         
-        # Create steps file
+        # Initialize SQLite storage (crash-safe, atomic)
+        import uuid
+        session_id = str(uuid.uuid4())[:8]
+        self.storage = EpiStorage(session_id, self.output_dir)
+        
+        # Keep JSONL path for backwards compatibility
         self.steps_file = self.output_dir / "steps.jsonl"
-        self.steps_file.touch()
+        
+        # Create empty steps.jsonl file immediately (for tests and early access)
+        self.steps_file.touch(exist_ok=True)
     
     def add_step(self, kind: str, content: Dict[str, Any]) -> None:
         """
@@ -93,24 +100,36 @@ class RecordingContext:
             f.write(step.model_dump_json() + '\n')
 
 
-# Global recording context (set by epi record command)
-_recording_context: Optional[RecordingContext] = None
+import contextvars
+
+# Thread-safe and async-safe recording context storage
+_recording_context: contextvars.ContextVar[Optional[RecordingContext]] = contextvars.ContextVar(
+    'epi_recording_context',
+    default=None
+)
 
 
-def set_recording_context(context: RecordingContext) -> None:
-    """Set global recording context."""
-    global _recording_context
-    _recording_context = context
+def set_recording_context(context: Optional[RecordingContext]) -> contextvars.Token:
+    """
+    Set recording context for current execution context (thread or async task).
+    
+    Args:
+        context: RecordingContext instance or None to clear
+        
+    Returns:
+        Token for resetting context later
+    """
+    return _recording_context.set(context)
 
 
 def get_recording_context() -> Optional[RecordingContext]:
-    """Get global recording context."""
-    return _recording_context
+    """Get recording context for current execution context."""
+    return _recording_context.get()
 
 
 def is_recording() -> bool:
-    """Check if recording is active."""
-    return _recording_context is not None
+    """Check if recording is active in current execution context."""
+    return _recording_context.get() is not None
 
 
 # ==================== OpenAI Patcher ====================
@@ -155,8 +174,9 @@ def _patch_openai_v1() -> bool:
         from openai import OpenAI
         from openai.resources.chat import completions
         
-        # Store original method
+        # Store original method for unpatching
         original_create = completions.Completions.create
+        _original_methods["openai.chat.completions.create"] = original_create
         
         @wraps(original_create)
         def wrapped_create(self, *args, **kwargs):
@@ -534,13 +554,53 @@ def patch_all() -> Dict[str, bool]:
     return results
 
 
+# Store original methods for unpatching
+_original_methods: Dict[str, Any] = {}
+
+
 def unpatch_all() -> None:
     """
     Unpatch all providers (restore original methods).
     
-    Note: This is a placeholder for future implementation.
-    Full unpatching requires storing original methods.
+    Restores any methods that were patched by patch_all(), patch_openai(),
+    patch_gemini(), or patch_requests().
     """
-    # For MVP, we don't implement unpatching
-    # In production, store original methods and restore them
-    pass
+    global _original_methods
+    
+    # Restore OpenAI v1+ if patched
+    if "openai.chat.completions.create" in _original_methods:
+        try:
+            from openai.resources.chat import completions
+            completions.Completions.create = _original_methods["openai.chat.completions.create"]
+        except ImportError:
+            pass
+    
+    # Restore OpenAI legacy if patched
+    if "openai.ChatCompletion.create" in _original_methods:
+        try:
+            import openai
+            openai.ChatCompletion.create = _original_methods["openai.ChatCompletion.create"]
+        except (ImportError, AttributeError):
+            pass
+    
+    # Restore Gemini if patched
+    if "gemini.generate_content" in _original_methods:
+        try:
+            import google.generativeai as genai
+            genai.GenerativeModel.generate_content = _original_methods["gemini.generate_content"]
+        except ImportError:
+            pass
+    
+    # Restore requests if patched
+    if "requests.Session.request" in _original_methods:
+        try:
+            import requests
+            requests.Session.request = _original_methods["requests.Session.request"]
+        except ImportError:
+            pass
+    
+    # Clear stored originals
+    _original_methods.clear()
+
+
+ 

epi_recorder/test_import.py

@@ -16,3 +16,5 @@ try:
     print("epi_cli imported successfully")
 except ImportError as e:
     print(f"Failed to import epi_cli: {e}")
+
+

epi_recorder/test_script.py

@@ -2,3 +2,5 @@
 print("Hello from EPI Test Script")
 import sys
 print(f"Python version: {sys.version}")
+
+

epi_recorder/wrappers/__init__.py

@@ -0,0 +1,16 @@
+"""
+EPI Wrapper Clients - Proxy wrappers for LLM clients.
+
+Provides transparent tracing without monkey patching.
+"""
+
+from epi_recorder.wrappers.openai import wrap_openai, TracedOpenAI, TracedCompletions, TracedChat
+from epi_recorder.wrappers.base import TracedClientBase
+
+__all__ = [
+    "wrap_openai",
+    "TracedOpenAI",
+    "TracedCompletions",
+    "TracedChat",
+    "TracedClientBase",
+]

epi_recorder/wrappers/base.py

@@ -0,0 +1,79 @@
+"""
+Base classes for EPI traced clients.
+"""
+
+from abc import ABC, abstractmethod
+from typing import Any, Optional
+from datetime import datetime
+
+
+class TracedClientBase(ABC):
+    """
+    Base class for traced LLM client wrappers.
+    
+    Provides common functionality for logging LLM calls
+    to the active EPI recording session.
+    """
+    
+    def __init__(self, client: Any):
+        """
+        Initialize traced client wrapper.
+        
+        Args:
+            client: The original LLM client to wrap
+        """
+        self._client = client
+    
+    def _get_session(self):
+        """Get the current active EPI recording session."""
+        from epi_recorder.api import get_current_session
+        return get_current_session()
+    
+    def _log_request(self, provider: str, model: str, messages: list, **kwargs) -> None:
+        """Log an LLM request to the active session."""
+        session = self._get_session()
+        if session:
+            session.log_step("llm.request", {
+                "provider": provider,
+                "model": model,
+                "messages": messages,
+                "timestamp": datetime.utcnow().isoformat(),
+                **kwargs
+            })
+    
+    def _log_response(
+        self, 
+        provider: str, 
+        model: str, 
+        content: str,
+        usage: Optional[dict] = None,
+        latency_seconds: Optional[float] = None,
+        **kwargs
+    ) -> None:
+        """Log an LLM response to the active session."""
+        session = self._get_session()
+        if session:
+            response_data = {
+                "provider": provider,
+                "model": model,
+                "choices": [{"message": {"role": "assistant", "content": content}}],
+                "timestamp": datetime.utcnow().isoformat(),
+            }
+            if usage:
+                response_data["usage"] = usage
+            if latency_seconds is not None:
+                response_data["latency_seconds"] = round(latency_seconds, 3)
+            response_data.update(kwargs)
+            session.log_step("llm.response", response_data)
+    
+    def _log_error(self, provider: str, error: Exception, **kwargs) -> None:
+        """Log an LLM error to the active session."""
+        session = self._get_session()
+        if session:
+            session.log_step("llm.error", {
+                "provider": provider,
+                "error": str(error),
+                "error_type": type(error).__name__,
+                "timestamp": datetime.utcnow().isoformat(),
+                **kwargs
+            })

epi_recorder/wrappers/openai.py

@@ -0,0 +1,178 @@
+"""
+OpenAI wrapper for EPI tracing.
+
+Provides a proxy wrapper that automatically logs all LLM calls
+without monkey patching.
+"""
+
+import time
+from typing import Any, Optional
+from datetime import datetime
+
+from epi_recorder.wrappers.base import TracedClientBase
+
+
+class TracedCompletions:
+    """Proxy wrapper for openai.chat.completions."""
+    
+    def __init__(self, completions: Any, provider: str = "openai"):
+        self._completions = completions
+        self._provider = provider
+    
+    def _get_session(self):
+        """Get the current active EPI recording session."""
+        from epi_recorder.api import get_current_session
+        return get_current_session()
+    
+    def create(self, *args, **kwargs) -> Any:
+        """
+        Create a chat completion with automatic EPI tracing.
+        
+        All arguments are passed through to the underlying client.
+        """
+        session = self._get_session()
+        
+        # Extract request info
+        model = kwargs.get("model", "unknown")
+        messages = kwargs.get("messages", [])
+        
+        # Log request if session is active
+        if session:
+            session.log_step("llm.request", {
+                "provider": self._provider,
+                "model": model,
+                "messages": messages,
+                "timestamp": datetime.utcnow().isoformat(),
+            })
+        
+        # Call original method
+        start_time = time.time()
+        try:
+            response = self._completions.create(*args, **kwargs)
+            latency = time.time() - start_time
+            
+            # Log response if session is active
+            if session:
+                # Extract response content
+                choices = []
+                for choice in response.choices:
+                    msg = choice.message
+                    choices.append({
+                        "message": {
+                            "role": getattr(msg, "role", "assistant"),
+                            "content": getattr(msg, "content", ""),
+                        },
+                        "finish_reason": getattr(choice, "finish_reason", None),
+                    })
+                
+                # Extract usage
+                usage = None
+                if hasattr(response, "usage") and response.usage:
+                    usage = {
+                        "prompt_tokens": getattr(response.usage, "prompt_tokens", 0),
+                        "completion_tokens": getattr(response.usage, "completion_tokens", 0),
+                        "total_tokens": getattr(response.usage, "total_tokens", 0),
+                    }
+                
+                session.log_step("llm.response", {
+                    "provider": self._provider,
+                    "model": model,
+                    "choices": choices,
+                    "usage": usage,
+                    "latency_seconds": round(latency, 3),
+                    "timestamp": datetime.utcnow().isoformat(),
+                })
+            
+            return response
+            
+        except Exception as e:
+            latency = time.time() - start_time
+            
+            # Log error if session is active
+            if session:
+                session.log_step("llm.error", {
+                    "provider": self._provider,
+                    "model": model,
+                    "error": str(e),
+                    "error_type": type(e).__name__,
+                    "latency_seconds": round(latency, 3),
+                    "timestamp": datetime.utcnow().isoformat(),
+                })
+            
+            raise
+
+
+class TracedChat:
+    """Proxy wrapper for openai.chat."""
+    
+    def __init__(self, chat: Any, provider: str = "openai"):
+        self._chat = chat
+        self._provider = provider
+        self.completions = TracedCompletions(chat.completions, provider)
+
+
+class TracedOpenAI(TracedClientBase):
+    """
+    Traced OpenAI client wrapper.
+    
+    Wraps an OpenAI client and automatically logs all LLM calls
+    to the active EPI recording session.
+    
+    Usage:
+        from openai import OpenAI
+        from epi_recorder.wrappers import wrap_openai
+        
+        client = wrap_openai(OpenAI())
+        
+        with record("my_agent.epi"):
+            response = client.chat.completions.create(
+                model="gpt-4",
+                messages=[{"role": "user", "content": "Hello"}]
+            )
+    """
+    
+    def __init__(self, client: Any, provider: str = "openai"):
+        """
+        Initialize traced OpenAI client.
+        
+        Args:
+            client: OpenAI client instance
+            provider: Provider name for logging (default: "openai")
+        """
+        super().__init__(client)
+        self._provider = provider
+        self.chat = TracedChat(client.chat, provider)
+    
+    def __getattr__(self, name: str) -> Any:
+        """
+        Forward attribute access to underlying client.
+        
+        This allows access to non-chat APIs (embeddings, files, etc.)
+        without explicit wrapping.
+        """
+        return getattr(self._client, name)
+
+
+def wrap_openai(client: Any, provider: str = "openai") -> TracedOpenAI:
+    """
+    Wrap an OpenAI client for EPI tracing.
+    
+    Args:
+        client: OpenAI client instance
+        provider: Provider name for logging (default: "openai")
+        
+    Returns:
+        TracedOpenAI wrapper
+        
+    Usage:
+        from openai import OpenAI
+        from epi_recorder.wrappers import wrap_openai
+        
+        # Wrap the client once
+        client = wrap_openai(OpenAI())
+        
+        # Use normally - calls are automatically traced when inside record()
+        with record("my_agent.epi"):
+            response = client.chat.completions.create(...)
+    """
+    return TracedOpenAI(client, provider)