fleet-python 0.2.81__py3-none-any.whl → 0.2.83__py3-none-any.whl

fleet/_async/client.py

@@ -17,11 +17,17 @@
 import asyncio
 import base64
 import cloudpickle
+import dataclasses
 import httpx
 import json
 import logging
 import os
+from datetime import date, datetime
+from decimal import Decimal
+from enum import Enum
+from pathlib import Path
 from typing import List, Optional, Dict, Any, TYPE_CHECKING, Union
+from uuid import UUID
 
 from .base import EnvironmentBase, AsyncWrapper
 from ..models import (
@@ -37,12 +43,112 @@ from ..models import (
     TaskUpdateRequest,
     Run,
     HeartbeatResponse,
+    SessionIngestRequest,
+    SessionIngestMessage,
+    SessionIngestResponse,
+    SessionStatus,
+    JobSessionsResponse,
+    SessionTranscriptResponse,
 )
 from .tasks import Task
 
 if TYPE_CHECKING:
     from .verifiers import AsyncVerifierFunction
 
+
+def _json_default(x: Any) -> Any:
+    """Default JSON serializer for non-native types."""
+    if isinstance(x, (datetime, date)):
+        return x.isoformat()
+    if isinstance(x, (UUID, Path)):
+        return str(x)
+    if isinstance(x, Decimal):
+        return float(x)
+    if isinstance(x, Enum):
+        return x.value
+    if isinstance(x, bytes):
+        return base64.b64encode(x).decode("utf-8")
+    if isinstance(x, set):
+        return list(x)
+    if dataclasses.is_dataclass(x) and not isinstance(x, type):
+        return dataclasses.asdict(x)
+    # Handle objects with __dict__ (generic objects)
+    if hasattr(x, "__dict__"):
+        return x.__dict__
+    raise TypeError(f"Not JSON serializable: {type(x)}")
+
+
+def _to_dict(obj: Any) -> Any:
+    """Convert any object to a JSON-serializable dict/value.
+
+    Handles:
+    - Pydantic v2 models (model_dump)
+    - Pydantic v1 models (.dict())
+    - dataclasses (asdict)
+    - TypedDict (just dict at runtime)
+    - Objects with __dict__
+    - Primitives pass through
+    """
+    if obj is None:
+        return None
+
+    # Pydantic v2
+    if hasattr(obj, "model_dump"):
+        return obj.model_dump()
+
+    # Pydantic v1
+    if hasattr(obj, "dict") and callable(obj.dict):
+        return obj.dict()
+
+    # dataclass
+    if dataclasses.is_dataclass(obj) and not isinstance(obj, type):
+        return dataclasses.asdict(obj)
+
+    # Already a dict or list - recursively convert
+    if isinstance(obj, dict):
+        return {k: _to_dict(v) for k, v in obj.items()}
+    if isinstance(obj, list):
+        return [_to_dict(v) for v in obj]
+
+    # Primitives
+    if isinstance(obj, (str, int, float, bool, type(None))):
+        return obj
+
+    # bytes -> base64
+    if isinstance(obj, bytes):
+        return base64.b64encode(obj).decode("utf-8")
+
+    # datetime/date
+    if isinstance(obj, (datetime, date)):
+        return obj.isoformat()
+
+    # UUID, Path
+    if isinstance(obj, (UUID, Path)):
+        return str(obj)
+
+    # Enum
+    if isinstance(obj, Enum):
+        return obj.value
+
+    # Decimal
+    if isinstance(obj, Decimal):
+        return float(obj)
+
+    # set
+    if isinstance(obj, set):
+        return list(obj)
+
+    # Generic object with __dict__
+    if hasattr(obj, "__dict__"):
+        return {k: _to_dict(v) for k, v in obj.__dict__.items() if not k.startswith("_")}
+
+    # Fallback - try to convert, or return string representation
+    try:
+        json.dumps(obj)
+        return obj
+    except (TypeError, ValueError):
+        return str(obj)
+
 from .instance import (
     AsyncInstanceClient,
     ResetRequest,
@@ -70,6 +176,166 @@ from .resources.mcp import AsyncMCPResource
 logger = logging.getLogger(__name__)
 
 
+class AsyncSession:
+    """A session for logging agent interactions to Fleet.
+
+    This provides a simple interface for streaming messages during an agent run.
+    Messages are sent one-by-one as they happen.
+
+    Usage:
+        session = await fleet.session_async(
+            model="anthropic/claude-sonnet-4",
+            task_key="my_task",
+            instance_id=env.instance_id,
+        )
+
+        # Log messages as they happen
+        await session.log(history, response)
+
+        # Complete when done
+        await session.complete()  # or session.fail()
+    """
+
+    def __init__(
+        self,
+        client: "AsyncFleet",
+        session_id: Optional[str] = None,
+        job_id: Optional[str] = None,
+        config: Optional[Any] = None,
+        model: Optional[str] = None,
+        task_key: Optional[str] = None,
+        instance_id: Optional[str] = None,
+    ):
+        self.session_id = session_id
+        self.job_id = job_id
+        self.config = config
+        self.model = model
+        self.task_key = task_key
+        self.instance_id = instance_id
+        self._client = client
+        self._message_count = 0
+        self._logged_count = 0  # Track how many messages from history have been logged
+        self._config_sent = False  # Only send config/model/task_key/instance_id on first log
+
+    async def log(self, history: List[Any], response: Any) -> SessionIngestResponse:
+        """Log an LLM call to the session.
+
+        Pass the input history and the model response. The session tracks what's
+        already been logged and only sends new messages. Objects are automatically
+        serialized to JSON (supports Pydantic, dataclasses, TypedDict, etc.).
+
+        Example:
+            response = model.generate(history)
+            await session.log(history, response.content)
+
+        Args:
+            history: The input messages sent to the model
+            response: The model's response (any serializable object)
+
+        Returns:
+            SessionIngestResponse with updated message count
+        """
+        # Collect new history messages since last call
+        new_history = history[self._logged_count:]
+
+        # Update tracked count to include the response we're about to send
+        # This prevents the response from being sent again as "new history" in the next call
+        self._logged_count = len(history) + (1 if response is not None else 0)
+
+        # Build the payload - serialize history + response to JSON
+        payload: Dict[str, Any] = {
+            "history": [_to_dict(msg) for msg in new_history],
+            "response": _to_dict(response),
+        }
+        if self.session_id:
+            payload["session_id"] = self.session_id
+        if self.job_id:
+            payload["job_id"] = self.job_id
+        # Include config, model, task_key, instance_id on first log only
+        if not self._config_sent:
+            if self.config is not None:
+                payload["config"] = _to_dict(self.config)
+            if self.model is not None:
+                payload["model"] = self.model
+            if self.task_key is not None:
+                payload["task_key"] = self.task_key
+            if self.instance_id is not None:
+                payload["instance_id"] = self.instance_id
+            self._config_sent = True
+
+        if not new_history and response is None:
+            return SessionIngestResponse(
+                success=True,
+                session_id=self.session_id or "",
+                message_count=self._message_count,
+                created_new_session=False,
+            )
+
+        result = await self._client._ingest_raw(payload=payload)
+        self._message_count = result.message_count
+        # Update session_id if this was the first log (new session created)
+        if not self.session_id and result.session_id:
+            self.session_id = result.session_id
+        return result
+
+    async def complete(
+        self,
+        verifier_execution_id: Optional[str] = None,
+    ) -> SessionIngestResponse:
+        """Mark the session as completed successfully.
+
+        Args:
+            verifier_execution_id: Optional ID of the verifier execution record
+
+        Returns:
+            SessionIngestResponse with final state
+        """
+        from datetime import datetime
+
+        payload: Dict[str, Any] = {
+            "session_id": self.session_id,
+            "status": "completed",
+            "ended_at": datetime.now().isoformat(),
+        }
+        if verifier_execution_id:
+            payload["verifier_execution_id"] = verifier_execution_id
+
+        response = await self._client._ingest_raw(payload)
+        self._message_count = response.message_count
+        return response
+
+    async def fail(
+        self,
+        verifier_execution_id: Optional[str] = None,
+    ) -> SessionIngestResponse:
+        """Mark the session as failed.
+
+        Args:
+            verifier_execution_id: Optional ID of the verifier execution record
+
+        Returns:
+            SessionIngestResponse with final state
+        """
+        from datetime import datetime
+
+        payload: Dict[str, Any] = {
+            "session_id": self.session_id,
+            "status": "failed",
+            "ended_at": datetime.now().isoformat(),
+        }
+        if verifier_execution_id:
+            payload["verifier_execution_id"] = verifier_execution_id
+
+        response = await self._client._ingest_raw(payload)
+        self._message_count = response.message_count
+        return response
+
+    @property
+    def message_count(self) -> int:
+        """Get the current message count."""
+        return self._message_count
+
+
 class AsyncEnv(EnvironmentBase):
     def __init__(self, client: Optional[AsyncWrapper], **kwargs):
         super().__init__(**kwargs)
@@ -1025,6 +1291,289 @@ class AsyncFleet:
         )
         return TaskResponse(**response.json())
 
+    # Sessions API methods
+
+    async def list_job_sessions(self, job_id: str) -> JobSessionsResponse:
+        """List all sessions for a job, grouped by task.
+
+        Args:
+            job_id: The job ID
+
+        Returns:
+            JobSessionsResponse containing sessions grouped by task with statistics
+        """
+        response = await self.client.request("GET", f"/v1/sessions/job/{job_id}")
+        return JobSessionsResponse(**response.json())
+
+    async def get_session_transcript(self, session_id: str) -> SessionTranscriptResponse:
+        """Get the transcript for a specific session.
+
+        Args:
+            session_id: The session ID
+
+        Returns:
+            SessionTranscriptResponse containing task, instance, verifier result, and messages
+        """
+        response = await self.client.request(
+            "GET", f"/v1/sessions/{session_id}/transcript"
+        )
+        return SessionTranscriptResponse(**response.json())
+
+    async def _ingest(
+        self,
+        messages: List[Dict[str, Any]],
+        session_id: Optional[str] = None,
+        model: Optional[str] = None,
+        task_key: Optional[str] = None,
+        job_id: Optional[str] = None,
+        instance_id: Optional[str] = None,
+        status: Optional[str] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+        started_at: Optional[str] = None,
+        ended_at: Optional[str] = None,
+        verifier_execution_id: Optional[str] = None,
+    ) -> SessionIngestResponse:
+        """Internal method to ingest session data."""
+        message_objects = [SessionIngestMessage(**msg) for msg in messages]
+        request = SessionIngestRequest(
+            messages=message_objects,
+            session_id=session_id,
+            model=model,
+            task_key=task_key,
+            job_id=job_id,
+            instance_id=instance_id,
+            status=SessionStatus(status) if status else None,
+            metadata=metadata,
+            started_at=started_at,
+            ended_at=ended_at,
+            verifier_execution_id=verifier_execution_id,
+        )
+        response = await self.client.request(
+            "POST",
+            "/v1/sessions/ingest",
+            json=request.model_dump(exclude_none=True),
+        )
+        return SessionIngestResponse(**response.json())
+
+    async def _ingest_raw(
+        self,
+        payload: Dict[str, Any],
+    ) -> SessionIngestResponse:
+        """Internal method to ingest raw session data as JSON.
+
+        This sends the history and response as-is to the backend,
+        letting the backend handle format normalization.
+        """
+        # Pre-serialize with our custom handler to ensure all types are JSON-safe
+        json_str = json.dumps(payload, default=_json_default)
+        clean_payload = json.loads(json_str)
+
+        response = await self.client.request(
+            "POST",
+            "/v1/traces/logs",
+            json=clean_payload,
+        )
+        return SessionIngestResponse(**response.json())
+
+    def start_session(
+        self,
+        session_id: Optional[str] = None,
+        job_id: Optional[str] = None,
+        config: Optional[Any] = None,
+        model: Optional[str] = None,
+        task_key: Optional[str] = None,
+        instance_id: Optional[str] = None,
+    ) -> AsyncSession:
+        """Start a new session for logging agent interactions.
+
+        This returns a Session object. The session is created on the backend
+        when you call log() for the first time.
+
+        Args:
+            session_id: Optional existing session ID to resume
+            job_id: Optional job ID to associate with the session
+            config: Optional config object (e.g., GenerateContentConfig) to log
+            model: Optional model name to log
+            task_key: Optional Fleet task key
+            instance_id: Optional Fleet instance ID
+
+        Returns:
+            AsyncSession object with log(), complete(), and fail() methods
+
+        Example:
+            session = fleet_client.start_session(config=config, model="gpt-4", task_key="task_123")
+
+            # Log LLM calls during agent run
+            await session.log(history, response)
+
+            # Complete when done
+            await session.complete()
+        """
+        return AsyncSession(
+            client=self,
+            session_id=session_id,
+            job_id=job_id,
+            config=config,
+            model=model,
+            task_key=task_key,
+            instance_id=instance_id,
+        )
+
+    async def trace_job(self, name: str) -> str:
+        """Create a new trace job.
+
+        Args:
+            name: Name of the job
+
+        Returns:
+            The job_id string
+        """
+        from fleet.models import TraceJobRequest, TraceJobResponse
+
+        request = TraceJobRequest(name=name)
+        response = await self.client.request(
+            "POST",
+            "/v1/traces/jobs",
+            json=request.model_dump(),
+        )
+        result = TraceJobResponse(**response.json())
+        return result.job_id
+
+    async def create_session(
+        self,
+        model: Optional[str] = None,
+        task_key: Optional[str] = None,
+        job_id: Optional[str] = None,
+        instance_id: Optional[str] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+        started_at: Optional[str] = None,
+        initial_message: Optional[Dict[str, Any]] = None,
+    ) -> SessionIngestResponse:
+        """Create a new session, optionally with an initial message.
+
+        This is useful for streaming scenarios where you want to create
+        a session first and then append messages one by one.
+
+        Args:
+            model: Model identifier (e.g., "anthropic/claude-sonnet-4")
+            task_key: Task key to associate with the session
+            job_id: Job ID to associate with the session
+            instance_id: Instance ID to associate with the session
+            metadata: Additional metadata for the session
+            started_at: ISO timestamp when session started
+            initial_message: Optional first message dict with 'role' and 'content'
+
+        Returns:
+            SessionIngestResponse containing session_id
+
+        Example:
+            # Create session and get ID
+            session = await fleet.create_session(
+                model="anthropic/claude-sonnet-4",
+                task_key="my_task",
+                started_at=datetime.now().isoformat()
+            )
+            
+            # Append messages as they happen
+            await fleet.append_message(session.session_id, {"role": "user", "content": "Hello"})
+            await fleet.append_message(session.session_id, {"role": "assistant", "content": "Hi!"})
+        """
+        # Use a placeholder message if none provided
+        if initial_message:
+            messages = [initial_message]
+        else:
+            messages = [{"role": "system", "content": "[session created]"}]
+
+        return await self._ingest(
+            messages=messages,
+            model=model,
+            task_key=task_key,
+            job_id=job_id,
+            instance_id=instance_id,
+            status="running",
+            metadata=metadata,
+            started_at=started_at,
+        )
+
+    async def append_message(
+        self,
+        session_id: str,
+        message: Dict[str, Any],
+        status: Optional[str] = None,
+        ended_at: Optional[str] = None,
+    ) -> SessionIngestResponse:
+        """Append a single message to an existing session.
+
+        This is useful for streaming scenarios where you want to send
+        messages one by one as they happen.
+
+        Args:
+            session_id: The session ID to append to
+            message: Message dict with 'role' and 'content' keys.
+                Optional keys: 'tool_calls', 'tool_call_id', 'timestamp', 'tokens', 'metadata'
+            status: Optional status update ("running", "completed", "failed")
+            ended_at: ISO timestamp when session ended (set when completing)
+
+        Returns:
+            SessionIngestResponse with updated message count
+
+        Example:
+            # Append user message
+            await fleet.append_message(session_id, {"role": "user", "content": "What's 2+2?"})
+            
+            # Append assistant response
+            await fleet.append_message(session_id, {"role": "assistant", "content": "4"})
+            
+            # Complete the session
+            await fleet.append_message(
+                session_id,
+                {"role": "assistant", "content": "Done!"},
+                status="completed",
+                ended_at=datetime.now().isoformat()
+            )
+        """
+        return await self._ingest(
+            messages=[message],
+            session_id=session_id,
+            status=status,
+            ended_at=ended_at,
+        )
+
+    async def complete_session(
+        self,
+        session_id: str,
+        status: str = "completed",
+        ended_at: Optional[str] = None,
+        final_message: Optional[Dict[str, Any]] = None,
+    ) -> SessionIngestResponse:
+        """Mark a session as complete.
+
+        Args:
+            session_id: The session ID to complete
+            status: Final status ("completed", "failed", "cancelled")
+            ended_at: ISO timestamp when session ended (defaults to now)
+            final_message: Optional final message to append
+
+        Returns:
+            SessionIngestResponse with final state
+        """
+        from datetime import datetime as dt
+        
+        if ended_at is None:
+            ended_at = dt.now().isoformat()
+        
+        if final_message:
+            messages = [final_message]
+        else:
+            messages = [{"role": "system", "content": f"[session {status}]"}]
+
+        return await self._ingest(
+            messages=messages,
+            session_id=session_id,
+            status=status,
+            ended_at=ended_at,
+        )
+
     async def _create_verifier_from_data(
         self, verifier_id: str, verifier_key: str, verifier_code: str, verifier_sha: str, verifier_runtime_version: Optional[str] = None
     ) -> "AsyncVerifierFunction":