agentpool 2.2.3__py3-none-any.whl → 2.5.0__py3-none-any.whl

agentpool_storage/session_store.py

@@ -5,10 +5,6 @@ from __future__ import annotations
 from datetime import timedelta
 from typing import TYPE_CHECKING, Any, Self
 
-from sqlalchemy import delete, select
-from sqlalchemy.ext.asyncio import AsyncSession
-from sqlmodel import SQLModel
-
 from agentpool.log import get_logger
 from agentpool.sessions.models import SessionData
 from agentpool.utils.now import get_now
@@ -45,6 +41,8 @@ class SQLSessionStore:
 
     async def __aenter__(self) -> Self:
         """Initialize database connection and create tables."""
+        from sqlmodel import SQLModel
+
         self._engine = self._config.get_engine()
 
         async with self._engine.begin() as conn:
@@ -98,7 +96,6 @@ class SQLSessionStore:
             project_id=data.project_id,
             parent_id=data.parent_id,
             version=data.version,
-            title=data.title,
             cwd=data.cwd,
             created_at=data.created_at,
             last_active=data.last_active,
@@ -115,7 +112,6 @@ class SQLSessionStore:
             project_id=row.project_id,
             parent_id=row.parent_id,
             version=row.version,
-            title=row.title,
             cwd=row.cwd,
             created_at=row.created_at,
             last_active=row.last_active,
@@ -130,6 +126,9 @@ class SQLSessionStore:
         Args:
             data: Session data to persist
         """
+        from sqlalchemy import delete
+        from sqlalchemy.ext.asyncio import AsyncSession
+
         engine = self._get_engine()
 
         async with AsyncSession(engine) as session:
@@ -152,6 +151,9 @@ class SQLSessionStore:
         Returns:
             Session data if found, None otherwise
         """
+        from sqlalchemy import select
+        from sqlalchemy.ext.asyncio import AsyncSession
+
         engine = self._get_engine()
 
         async with AsyncSession(engine) as session:
@@ -173,6 +175,9 @@ class SQLSessionStore:
         Returns:
             True if session was deleted, False if not found
         """
+        from sqlalchemy import delete
+        from sqlalchemy.ext.asyncio import AsyncSession
+
         engine = self._get_engine()
 
         async with AsyncSession(engine) as session:
@@ -199,6 +204,9 @@ class SQLSessionStore:
         Returns:
             List of session IDs
         """
+        from sqlalchemy import select
+        from sqlalchemy.ext.asyncio import AsyncSession
+
         engine = self._get_engine()
 
         async with AsyncSession(engine) as session:
@@ -223,6 +231,9 @@ class SQLSessionStore:
         Returns:
             Number of sessions removed
         """
+        from sqlalchemy import delete
+        from sqlalchemy.ext.asyncio import AsyncSession
+
         engine = self._get_engine()
         cutoff = get_now() - timedelta(hours=max_age_hours)
 
@@ -250,6 +261,9 @@ class SQLSessionStore:
         Returns:
             List of session data objects
         """
+        from sqlalchemy import select
+        from sqlalchemy.ext.asyncio import AsyncSession
+
         engine = self._get_engine()
 
         async with AsyncSession(engine) as session:

agentpool_storage/sql_provider/sql_provider.py

@@ -130,7 +130,6 @@ class SQLModelProvider(StorageProvider):
         cost_info: TokenCost | None = None,
         model: str | None = None,
         response_time: float | None = None,
-        forwarded_from: list[str] | None = None,
         provider_name: str | None = None,
         provider_response_id: str | None = None,
         messages: str | None = None,
@@ -157,7 +156,6 @@ class SQLModelProvider(StorageProvider):
                 input_tokens=cost_info.token_usage.input_tokens if cost_info else None,
                 output_tokens=cost_info.token_usage.output_tokens if cost_info else None,
                 cost=float(cost_info.total_cost) if cost_info else None,
-                forwarded_from=forwarded_from,
                 provider_name=provider_name,
                 provider_response_id=provider_response_id,
                 messages=messages,
@@ -210,6 +208,141 @@ class SQLModelProvider(StorageProvider):
             )
             return result.scalar_one_or_none()
 
+    async def get_conversation_messages(
+        self,
+        conversation_id: str,
+        *,
+        include_ancestors: bool = False,
+    ) -> list[ChatMessage[str]]:
+        """Get all messages for a conversation.
+
+        Args:
+            conversation_id: ID of the conversation
+            include_ancestors: If True, traverse parent_id chain to include
+                messages from ancestor conversations (for forked convos).
+
+        Returns:
+            List of messages ordered by timestamp.
+        """
+        async with AsyncSession(self.engine) as session:
+            # Get messages for this conversation
+            result = await session.execute(
+                select(Message)
+                .where(Message.conversation_id == conversation_id)
+                .order_by(Message.timestamp.asc())  # type: ignore
+            )
+            messages = [to_chat_message(m) for m in result.scalars().all()]
+
+            if not include_ancestors or not messages:
+                return messages
+
+            # Find the first message's parent_id to get ancestor chain
+            first_msg = messages[0]
+            if first_msg.parent_id:
+                ancestors = await self.get_message_ancestry(first_msg.parent_id)
+                return ancestors + messages
+
+            return messages
+
+    async def get_message(
+        self,
+        message_id: str,
+    ) -> ChatMessage[str] | None:
+        """Get a single message by ID."""
+        async with AsyncSession(self.engine) as session:
+            result = await session.execute(select(Message).where(Message.id == message_id))
+            msg = result.scalar_one_or_none()
+            return to_chat_message(msg) if msg else None
+
+    async def get_message_ancestry(
+        self,
+        message_id: str,
+    ) -> list[ChatMessage[str]]:
+        """Get the ancestry chain of a message.
+
+        Traverses parent_id chain to build full history.
+        """
+        ancestors: list[ChatMessage[str]] = []
+        current_id: str | None = message_id
+
+        async with AsyncSession(self.engine) as session:
+            while current_id:
+                result = await session.execute(select(Message).where(Message.id == current_id))
+                msg = result.scalar_one_or_none()
+                if not msg:
+                    break
+                ancestors.append(to_chat_message(msg))
+                current_id = msg.parent_id
+
+        # Reverse to get oldest first
+        ancestors.reverse()
+        return ancestors
+
+    async def fork_conversation(
+        self,
+        *,
+        source_conversation_id: str,
+        new_conversation_id: str,
+        fork_from_message_id: str | None = None,
+        new_agent_name: str | None = None,
+    ) -> str | None:
+        """Fork a conversation at a specific point.
+
+        Creates a new conversation record. The fork point message_id is returned
+        so callers can set it as parent_id for new messages.
+        """
+        async with AsyncSession(self.engine) as session:
+            # Get source conversation
+            result = await session.execute(
+                select(Conversation).where(Conversation.id == source_conversation_id)
+            )
+            source_conv = result.scalar_one_or_none()
+            if not source_conv:
+                msg = f"Source conversation not found: {source_conversation_id}"
+                raise ValueError(msg)
+
+            # Determine fork point
+            fork_point_id: str | None = None
+            if fork_from_message_id:
+                # Verify the message exists and belongs to the source conversation
+                msg_result = await session.execute(
+                    select(Message).where(
+                        Message.id == fork_from_message_id,
+                        Message.conversation_id == source_conversation_id,
+                    )
+                )
+                fork_msg = msg_result.scalar_one_or_none()
+                if not fork_msg:
+                    err = f"Message {fork_from_message_id} not found in conversation"
+                    raise ValueError(err)
+                fork_point_id = fork_from_message_id
+            else:
+                # Fork from the last message
+                msg_result = await session.execute(
+                    select(Message)
+                    .where(Message.conversation_id == source_conversation_id)
+                    .order_by(desc(Message.timestamp))
+                    .limit(1)
+                )
+                last_msg = msg_result.scalar_one_or_none()
+                if last_msg:
+                    fork_point_id = last_msg.id
+
+            # Create new conversation
+            agent_name = new_agent_name or source_conv.agent_name
+            new_conv = Conversation(
+                id=new_conversation_id,
+                agent_name=agent_name,
+                title=f"{source_conv.title or 'Conversation'} (fork)"
+                if source_conv.title
+                else None,
+                start_time=get_now(),
+            )
+            session.add(new_conv)
+            await session.commit()
+
+            return fork_point_id
+
     async def log_command(
         self,
         *,

agentpool_storage/sql_provider/utils.py

@@ -8,8 +8,7 @@ from decimal import Decimal
 from typing import TYPE_CHECKING, Any
 
 from pydantic_ai import RunUsage
-from sqlalchemy import JSON, Column, and_, or_
-from sqlalchemy.sql import expression
+from sqlalchemy import Column, and_
 from sqlmodel import select
 
 from agentpool.messaging import ChatMessage, TokenCost
@@ -68,7 +67,6 @@ def to_chat_message(db_message: Message) -> ChatMessage[str]:
         model_name=db_message.model,
         cost_info=cost_info,
         response_time=db_message.response_time,
-        forwarded_from=db_message.forwarded_from or [],
         timestamp=db_message.timestamp,
         provider_name=db_message.provider_name,
         provider_response_id=db_message.provider_response_id,
@@ -168,15 +166,7 @@ def build_message_query(query: SessionQuery) -> SelectOfScalar[Any]:
     if query.name:
         conditions.append(Message.conversation_id == query.name)
     if query.agents:
-        agent_conditions = [Column("name").in_(query.agents)]
-        if query.include_forwarded:
-            agent_conditions.append(
-                and_(
-                    Column("forwarded_from").isnot(None),
-                    expression.cast(Column("forwarded_from"), JSON).contains(list(query.agents)),  # type: ignore
-                )
-            )
-        conditions.append(or_(*agent_conditions))
+        conditions.append(Column("name").in_(query.agents))
     if query.since and (cutoff := query.get_time_cutoff()):
         conditions.append(Message.timestamp >= cutoff)
     if query.until:

agentpool_storage/zed_provider/__init__.py

@@ -0,0 +1,16 @@
+"""Zed IDE storage provider.
+
+This package implements a read-only storage backend that reads Zed IDE's
+native thread format from ~/.local/share/zed/threads/threads.db.
+
+Zed stores conversations as zstd-compressed JSON in a SQLite database.
+This provider enables importing and analyzing Zed conversations.
+"""
+
+from __future__ import annotations
+
+from agentpool_storage.zed_provider.provider import ZedStorageProvider
+
+__all__ = [
+    "ZedStorageProvider",
+]

agentpool_storage/zed_provider/helpers.py

@@ -0,0 +1,281 @@
+"""Helper functions for Zed storage provider.
+
+Stateless conversion and utility functions for working with Zed format.
+"""
+
+from __future__ import annotations
+
+import base64
+from datetime import datetime
+import io
+import json
+from typing import Any
+
+from pydantic_ai.messages import (
+    BinaryContent,
+    ModelRequest,
+    ModelResponse,
+    TextPart,
+    ThinkingPart,
+    ToolCallPart,
+    ToolReturnPart,
+    UserPromptPart,
+)
+from pydantic_ai.usage import RequestUsage
+import zstandard
+
+from agentpool.log import get_logger
+from agentpool.messaging import ChatMessage
+from agentpool.utils.now import get_now
+from agentpool_storage.zed_provider.models import ZedThread
+
+
+logger = get_logger(__name__)
+
+
+def decompress_thread(data: bytes, data_type: str) -> ZedThread:
+    """Decompress and parse thread data.
+
+    Args:
+        data: Compressed thread data
+        data_type: Type of compression ("zstd" or plain)
+
+    Returns:
+        Parsed ZedThread object
+    """
+    if data_type == "zstd":
+        dctx = zstandard.ZstdDecompressor()
+        # Use stream_reader for data without content size in header
+        reader = dctx.stream_reader(io.BytesIO(data))
+        json_data = reader.read()
+    else:
+        json_data = data
+
+    thread_dict = json.loads(json_data)
+    return ZedThread.model_validate(thread_dict)
+
+
+def parse_user_content(  # noqa: PLR0915
+    content_list: list[dict[str, Any]],
+) -> tuple[str, list[str | BinaryContent]]:
+    """Parse user message content blocks.
+
+    Args:
+        content_list: List of content blocks from Zed user message
+
+    Returns:
+        Tuple of (display_text, pydantic_ai_content_list)
+    """
+    display_parts: list[str] = []
+    pydantic_content: list[str | BinaryContent] = []
+
+    for item in content_list:
+        if "Text" in item:
+            text = item["Text"]
+            display_parts.append(text)
+            pydantic_content.append(text)
+
+        elif "Image" in item:
+            image_data = item["Image"]
+            source = image_data.get("source", "")
+            try:
+                binary_data = base64.b64decode(source)
+                # Try to detect image type from magic bytes
+                media_type = "image/png"  # default
+                if binary_data[:3] == b"\xff\xd8\xff":
+                    media_type = "image/jpeg"
+                elif binary_data[:4] == b"\x89PNG":
+                    media_type = "image/png"
+                elif binary_data[:6] in (b"GIF87a", b"GIF89a"):
+                    media_type = "image/gif"
+                elif binary_data[:4] == b"RIFF" and binary_data[8:12] == b"WEBP":
+                    media_type = "image/webp"
+
+                pydantic_content.append(BinaryContent(data=binary_data, media_type=media_type))
+                display_parts.append("[image]")
+            except (ValueError, TypeError, IndexError) as e:
+                logger.warning("Failed to decode image", error=str(e))
+                display_parts.append("[image decode error]")
+
+        elif "Mention" in item:
+            mention = item["Mention"]
+            uri = mention.get("uri", {})
+            content = mention.get("content", "")
+
+            # Format mention based on type
+            if "File" in uri:
+                file_info = uri["File"]
+                path = file_info.get("abs_path", "unknown")
+                formatted = f"[File: {path}]\n{content}"
+            elif "Directory" in uri:
+                dir_info = uri["Directory"]
+                path = dir_info.get("abs_path", "unknown")
+                formatted = f"[Directory: {path}]\n{content}"
+            elif "Symbol" in uri:
+                symbol_info = uri["Symbol"]
+                path = symbol_info.get("abs_path", "unknown")
+                name = symbol_info.get("name", "")
+                formatted = f"[Symbol: {name} in {path}]\n{content}"
+            elif "Selection" in uri:
+                sel_info = uri["Selection"]
+                path = sel_info.get("abs_path", "unknown")
+                formatted = f"[Selection: {path}]\n{content}"
+            elif "Fetch" in uri:
+                fetch_info = uri["Fetch"]
+                url = fetch_info.get("url", "unknown")
+                formatted = f"[Fetched: {url}]\n{content}"
+            else:
+                formatted = content
+
+            display_parts.append(formatted)
+            pydantic_content.append(formatted)
+
+    display_text = "\n".join(display_parts)
+    return display_text, pydantic_content
+
+
+def parse_agent_content(
+    content_list: list[dict[str, Any]],
+) -> tuple[str, list[TextPart | ThinkingPart | ToolCallPart]]:
+    """Parse agent message content blocks.
+
+    Args:
+        content_list: List of content blocks from Zed agent message
+
+    Returns:
+        Tuple of (display_text, pydantic_ai_parts)
+    """
+    display_parts: list[str] = []
+    pydantic_parts: list[TextPart | ThinkingPart | ToolCallPart] = []
+
+    for item in content_list:
+        if "Text" in item:
+            text = item["Text"]
+            display_parts.append(text)
+            pydantic_parts.append(TextPart(content=text))
+
+        elif "Thinking" in item:
+            thinking = item["Thinking"]
+            text = thinking.get("text", "")
+            signature = thinking.get("signature")
+            display_parts.append(f"<thinking>\n{text}\n</thinking>")
+            pydantic_parts.append(ThinkingPart(content=text, signature=signature))
+
+        elif "ToolUse" in item:
+            tool_use = item["ToolUse"]
+            tool_id = tool_use.get("id", "")
+            tool_name = tool_use.get("name", "")
+            tool_input = tool_use.get("input", {})
+            display_parts.append(f"[Tool: {tool_name}]")
+            pydantic_parts.append(
+                ToolCallPart(tool_name=tool_name, args=tool_input, tool_call_id=tool_id)
+            )
+
+    display_text = "\n".join(display_parts)
+    return display_text, pydantic_parts
+
+
+def parse_tool_results(tool_results: dict[str, Any]) -> list[ToolReturnPart]:
+    """Parse tool results into ToolReturnParts.
+
+    Args:
+        tool_results: Dictionary of tool results from Zed agent message
+
+    Returns:
+        List of ToolReturnPart objects
+    """
+    parts: list[ToolReturnPart] = []
+
+    for tool_id, result in tool_results.items():
+        if isinstance(result, dict):
+            tool_name = result.get("tool_name", "")
+            output = result.get("output", "")
+            content = result.get("content", {})
+
+            # Handle output being a dict like {"Text": "..."}
+            if isinstance(output, dict) and "Text" in output:
+                output = output["Text"]
+            # Extract text content if available
+            elif isinstance(content, dict) and "Text" in content:
+                output = content["Text"]
+            elif isinstance(content, str):
+                output = content
+
+            parts.append(
+                ToolReturnPart(tool_name=tool_name, content=output or "", tool_call_id=tool_id)
+            )
+
+    return parts
+
+
+def thread_to_chat_messages(thread: ZedThread, thread_id: str) -> list[ChatMessage[str]]:
+    """Convert a Zed thread to ChatMessages.
+
+    Args:
+        thread: Zed thread object
+        thread_id: Thread identifier
+
+    Returns:
+        List of ChatMessage objects
+    """
+    messages: list[ChatMessage[str]] = []
+    try:
+        updated_at = datetime.fromisoformat(thread.updated_at.replace("Z", "+00:00"))
+    except (ValueError, AttributeError):
+        updated_at = get_now()
+    # Get model info
+    model_name = None
+    if thread.model:
+        model_name = f"{thread.model.provider}:{thread.model.model}"
+
+    for idx, msg in enumerate(thread.messages):
+        if msg == "Resume":
+            continue  # Skip control messages
+        msg_id = f"{thread_id}_{idx}"
+
+        if msg.User is not None:
+            user_msg = msg.User
+            display_text, pydantic_content = parse_user_content(user_msg.content)
+            part = UserPromptPart(content=pydantic_content)
+            model_request = ModelRequest(parts=[part])
+
+            messages.append(
+                ChatMessage[str](
+                    content=display_text,
+                    conversation_id=thread_id,
+                    role="user",
+                    message_id=user_msg.id or msg_id,
+                    name=None,
+                    model_name=None,
+                    timestamp=updated_at,  # Zed doesn't store per-message timestamps
+                    messages=[model_request],
+                )
+            )
+
+        elif msg.Agent is not None:
+            agent_msg = msg.Agent
+            display_text, pydantic_parts = parse_agent_content(agent_msg.content)
+            # Build ModelResponse
+            usage = RequestUsage()
+            model_response = ModelResponse(parts=pydantic_parts, usage=usage, model_name=model_name)
+            # Build tool return parts for next request if there are tool results
+            tool_return_parts = parse_tool_results(agent_msg.tool_results)
+            # Create the messages list - response, then optionally tool returns
+            pydantic_messages: list[ModelResponse | ModelRequest] = [model_response]
+            if tool_return_parts:
+                pydantic_messages.append(ModelRequest(parts=tool_return_parts))
+
+            messages.append(
+                ChatMessage[str](
+                    content=display_text,
+                    conversation_id=thread_id,
+                    role="assistant",
+                    message_id=msg_id,
+                    name="zed",
+                    model_name=model_name,
+                    timestamp=updated_at,
+                    messages=pydantic_messages,
+                )
+            )
+
+    return messages