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

agentpool_storage/file_provider.py

@@ -4,9 +4,9 @@ from __future__ import annotations
 
 from datetime import datetime
 from decimal import Decimal
-from typing import TYPE_CHECKING, TypedDict, cast
+from typing import TYPE_CHECKING, Any, TypedDict, cast
 
-from pydantic_ai import RunUsage
+from pydantic_ai import FinishReason, RunUsage  # noqa: TC002
 from upathtools import to_upath
 
 from agentpool.common_types import JsonValue, MessageRole  # noqa: TC001
@@ -19,9 +19,9 @@ from agentpool_storage.models import TokenUsage
 
 
 if TYPE_CHECKING:
-    from pydantic_ai import FinishReason
     from yamling import FormatType
 
+    from agentpool.sessions.models import ProjectData
     from agentpool_config.session import SessionQuery
     from agentpool_config.storage import FileStorageConfig
 
@@ -41,11 +41,11 @@ class MessageData(TypedDict):
     cost: Decimal | None
     token_usage: TokenUsage | None
     response_time: float | None
-    forwarded_from: list[str] | None
     provider_name: str | None
     provider_response_id: str | None
     messages: str | None
     finish_reason: FinishReason | None
+    parent_id: str | None
 
 
 class ConversationData(TypedDict):
@@ -68,12 +68,26 @@ class CommandData(TypedDict):
     metadata: dict[str, JsonValue]
 
 
+class ProjectDataDict(TypedDict):
+    """Data structure for storing project information."""
+
+    project_id: str
+    worktree: str
+    name: str | None
+    vcs: str | None
+    config_path: str | None
+    created_at: str
+    last_active: str
+    settings: dict[str, JsonValue]
+
+
 class StorageData(TypedDict):
     """Data structure for storing storage information."""
 
     messages: list[MessageData]
     conversations: list[ConversationData]
     commands: list[CommandData]
+    projects: list[ProjectDataDict]
 
 
 class FileProvider(StorageProvider):
@@ -99,6 +113,7 @@ class FileProvider(StorageProvider):
             "messages": [],
             "conversations": [],
             "commands": [],
+            "projects": [],
         }
         self._load()
 
@@ -132,14 +147,7 @@ class FileProvider(StorageProvider):
             # Apply filters
             if query.name and msg["conversation_id"] != query.name:
                 continue
-            if query.agents and not (
-                msg["name"] in query.agents
-                or (
-                    query.include_forwarded
-                    and msg["forwarded_from"]
-                    and any(a in query.agents for a in msg["forwarded_from"])
-                )
-            ):
+            if query.agents and msg["name"] not in query.agents:
                 continue
             cutoff = query.get_time_cutoff()
             timestamp = datetime.fromisoformat(msg["timestamp"])
@@ -175,7 +183,6 @@ class FileProvider(StorageProvider):
                 model_name=msg["model"],
                 cost_info=cost_info,
                 response_time=msg["response_time"],
-                forwarded_from=msg["forwarded_from"] or [],
                 timestamp=datetime.fromisoformat(msg["timestamp"]),
                 provider_name=msg["provider_name"],
                 provider_response_id=msg["provider_response_id"],
@@ -200,11 +207,11 @@ class FileProvider(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,
         finish_reason: FinishReason | None = None,
+        parent_id: str | None = None,
     ) -> None:
         """Log a new message."""
         self._data["messages"].append({
@@ -222,11 +229,11 @@ class FileProvider(StorageProvider):
                 total=cost_info.token_usage.total_tokens if cost_info else 0,
             ),
             "response_time": response_time,
-            "forwarded_from": forwarded_from,
             "provider_name": provider_name,
             "provider_response_id": provider_response_id,
             "messages": messages,
             "finish_reason": finish_reason,
+            "parent_id": parent_id,
         })
         self._save()
 
@@ -269,6 +276,152 @@ class FileProvider(StorageProvider):
                 return conv.get("title")
         return None
 
+    async def get_conversation_messages(
+        self,
+        conversation_id: str,
+        *,
+        include_ancestors: bool = False,
+    ) -> list[ChatMessage[str]]:
+        """Get all messages for a conversation."""
+        messages: list[ChatMessage[str]] = []
+        for msg in self._data["messages"]:
+            if msg["conversation_id"] != conversation_id:
+                continue
+            chat_msg = self._to_chat_message(msg)
+            messages.append(chat_msg)
+
+        # Sort by timestamp
+        messages.sort(key=lambda m: m.timestamp or get_now())
+
+        if not include_ancestors or not messages:
+            return messages
+
+        # Get ancestor chain if first message has parent_id
+        first_msg = messages[0]
+        if first_msg.parent_id:
+            ancestors = await self.get_message_ancestry(first_msg.parent_id)
+            return ancestors + messages
+
+        return messages
+
+    def _to_chat_message(self, msg: MessageData) -> ChatMessage[str]:
+        """Convert stored message data to ChatMessage."""
+        cost_info = None
+        if msg.get("token_usage"):
+            usage = msg["token_usage"]
+            cost_info = TokenCost(
+                token_usage=RunUsage(
+                    input_tokens=usage.get("prompt", 0) if usage else 0,
+                    output_tokens=usage.get("completion", 0) if usage else 0,
+                ),
+                total_cost=Decimal(str(msg.get("cost") or 0)),
+            )
+
+        # Build kwargs, only including timestamp/message_id if they have values
+        kwargs: dict[str, Any] = {
+            "content": msg["content"],
+            "role": cast(MessageRole, msg["role"]),
+            "name": msg.get("name"),
+            "model_name": msg.get("model"),
+            "cost_info": cost_info,
+            "response_time": msg.get("response_time"),
+            "parent_id": msg.get("parent_id"),
+            "conversation_id": msg.get("conversation_id"),
+            "messages": deserialize_messages(msg.get("messages")),
+            "finish_reason": msg.get("finish_reason"),
+        }
+        if msg.get("timestamp"):
+            kwargs["timestamp"] = datetime.fromisoformat(msg["timestamp"])
+        if msg.get("message_id"):
+            kwargs["message_id"] = msg["message_id"]
+
+        return ChatMessage[str](**kwargs)
+
+    async def get_message(self, message_id: str) -> ChatMessage[str] | None:
+        """Get a single message by ID."""
+        for msg in self._data["messages"]:
+            if msg.get("message_id") == message_id:
+                return self._to_chat_message(msg)
+        return None
+
+    async def get_message_ancestry(self, message_id: str) -> list[ChatMessage[str]]:
+        """Get the ancestry chain of a message."""
+        ancestors: list[ChatMessage[str]] = []
+        current_id: str | None = message_id
+
+        while current_id:
+            msg = await self.get_message(current_id)
+            if not msg:
+                break
+            ancestors.append(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."""
+        # Find source conversation
+        source_conv = next(
+            (c for c in self._data["conversations"] if c["id"] == source_conversation_id),
+            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 message exists in source conversation
+            msg_exists = any(
+                m.get("message_id") == fork_from_message_id
+                and m["conversation_id"] == source_conversation_id
+                for m in self._data["messages"]
+            )
+            if not msg_exists:
+                err = f"Message {fork_from_message_id} not found in conversation"
+                raise ValueError(err)
+            fork_point_id = fork_from_message_id
+        else:
+            # Find last message in source conversation
+            conv_messages = [
+                m for m in self._data["messages"] if m["conversation_id"] == source_conversation_id
+            ]
+            if conv_messages:
+                conv_messages.sort(
+                    key=lambda m: (
+                        datetime.fromisoformat(m["timestamp"]) if m.get("timestamp") else get_now()
+                    )
+                )
+                fork_point_id = conv_messages[-1].get("message_id")
+
+        # Create new conversation
+        agent_name = new_agent_name or source_conv["agent_name"]
+        title = (
+            f"{source_conv.get('title') or 'Conversation'} (fork)"
+            if source_conv.get("title")
+            else None
+        )
+        new_conv: ConversationData = {
+            "id": new_conversation_id,
+            "agent_name": agent_name,
+            "title": title,
+            "start_time": get_now().isoformat(),
+        }
+        self._data["conversations"].append(new_conv)
+        self._save()
+
+        return fork_point_id
+
     async def log_command(
         self,
         *,
@@ -329,6 +482,7 @@ class FileProvider(StorageProvider):
                 "messages": [],
                 "conversations": [],
                 "commands": [],
+                "projects": [],
             }
             self._save()
             return conv_count, msg_count
@@ -376,3 +530,112 @@ class FileProvider(StorageProvider):
             msg_count = len(self._data["messages"])
 
         return conv_count, msg_count
+
+    async def delete_conversation_messages(
+        self,
+        conversation_id: str,
+    ) -> int:
+        """Delete all messages for a conversation."""
+        original_count = len(self._data["messages"])
+        self._data["messages"] = [
+            m for m in self._data["messages"] if m["conversation_id"] != conversation_id
+        ]
+        deleted = original_count - len(self._data["messages"])
+        if deleted > 0:
+            self._save()
+        return deleted
+
+    # Project methods
+
+    def _to_project_data(self, data: ProjectDataDict) -> ProjectData:
+        """Convert dict to ProjectData."""
+        from datetime import datetime
+
+        from agentpool.sessions.models import ProjectData
+
+        return ProjectData(
+            project_id=data["project_id"],
+            worktree=data["worktree"],
+            name=data["name"],
+            vcs=data["vcs"],
+            config_path=data["config_path"],
+            created_at=datetime.fromisoformat(data["created_at"]),
+            last_active=datetime.fromisoformat(data["last_active"]),
+            settings=data["settings"],
+        )
+
+    def _to_project_dict(self, project: ProjectData) -> ProjectDataDict:
+        """Convert ProjectData to dict."""
+        return ProjectDataDict(
+            project_id=project.project_id,
+            worktree=project.worktree,
+            name=project.name,
+            vcs=project.vcs,
+            config_path=project.config_path,
+            created_at=project.created_at.isoformat(),
+            last_active=project.last_active.isoformat(),
+            settings=project.settings,
+        )
+
+    async def save_project(self, project: ProjectData) -> None:
+        """Save or update a project."""
+        # Remove existing if present
+        self._data["projects"] = [
+            p for p in self._data.get("projects", []) if p["project_id"] != project.project_id
+        ]
+        # Add new/updated
+        self._data["projects"].append(self._to_project_dict(project))
+        self._save()
+        logger.debug("Saved project", project_id=project.project_id)
+
+    async def get_project(self, project_id: str) -> ProjectData | None:
+        """Get a project by ID."""
+        for p in self._data.get("projects", []):
+            if p["project_id"] == project_id:
+                return self._to_project_data(p)
+        return None
+
+    async def get_project_by_worktree(self, worktree: str) -> ProjectData | None:
+        """Get a project by worktree path."""
+        for p in self._data.get("projects", []):
+            if p["worktree"] == worktree:
+                return self._to_project_data(p)
+        return None
+
+    async def get_project_by_name(self, name: str) -> ProjectData | None:
+        """Get a project by friendly name."""
+        for p in self._data.get("projects", []):
+            if p["name"] == name:
+                return self._to_project_data(p)
+        return None
+
+    async def list_projects(self, limit: int | None = None) -> list[ProjectData]:
+        """List all projects, ordered by last_active descending."""
+        projects = sorted(
+            self._data.get("projects", []),
+            key=lambda p: p["last_active"],
+            reverse=True,
+        )
+        if limit is not None:
+            projects = projects[:limit]
+        return [self._to_project_data(p) for p in projects]
+
+    async def delete_project(self, project_id: str) -> bool:
+        """Delete a project."""
+        original_count = len(self._data.get("projects", []))
+        self._data["projects"] = [
+            p for p in self._data.get("projects", []) if p["project_id"] != project_id
+        ]
+        deleted = original_count > len(self._data["projects"])
+        if deleted:
+            self._save()
+            logger.debug("Deleted project", project_id=project_id)
+        return deleted
+
+    async def touch_project(self, project_id: str) -> None:
+        """Update project's last_active timestamp."""
+        for p in self._data.get("projects", []):
+            if p["project_id"] == project_id:
+                p["last_active"] = get_now().isoformat()
+                self._save()
+                return

agentpool_storage/memory_provider.py

@@ -16,11 +16,37 @@ if TYPE_CHECKING:
     from collections.abc import Sequence
 
     from agentpool.common_types import JsonValue
+    from agentpool.sessions.models import ProjectData
     from agentpool_config.session import SessionQuery
     from agentpool_config.storage import MemoryStorageConfig
     from agentpool_storage.models import MessageData, QueryFilters, StatsFilters, TokenUsage
 
 
+def _dict_to_chat_message(msg: dict[str, Any]) -> ChatMessage[str]:
+    """Convert a stored message dict to ChatMessage."""
+    cost_info = None
+    if msg.get("cost_info"):
+        cost_info = TokenCost(token_usage=msg["cost_info"], total_cost=msg.get("cost", 0.0))
+
+    # Build kwargs, only including timestamp/message_id if they exist
+    kwargs: dict[str, Any] = {
+        "content": msg["content"],
+        "role": msg["role"],
+        "name": msg.get("name"),
+        "model_name": msg.get("model"),
+        "cost_info": cost_info,
+        "response_time": msg.get("response_time"),
+        "parent_id": msg.get("parent_id"),
+        "conversation_id": msg.get("conversation_id"),
+    }
+    if msg.get("timestamp"):
+        kwargs["timestamp"] = msg["timestamp"]
+    if msg.get("message_id"):
+        kwargs["message_id"] = msg["message_id"]
+
+    return ChatMessage[str](**kwargs)
+
+
 class MemoryStorageProvider(StorageProvider):
     """In-memory storage provider for testing."""
 
@@ -31,12 +57,14 @@ class MemoryStorageProvider(StorageProvider):
         self.messages: list[dict[str, Any]] = []
         self.conversations: list[dict[str, Any]] = []
         self.commands: list[dict[str, Any]] = []
+        self.projects: dict[str, ProjectData] = {}
 
     def cleanup(self) -> None:
         """Clear all stored data."""
         self.messages.clear()
         self.conversations.clear()
         self.commands.clear()
+        self.projects.clear()
 
     async def filter_messages(self, query: SessionQuery) -> list[ChatMessage[str]]:
         """Filter messages from memory."""
@@ -49,14 +77,7 @@ class MemoryStorageProvider(StorageProvider):
                 continue
 
             # Skip if agent name doesn't match
-            if query.agents and not (
-                msg["name"] in query.agents
-                or (
-                    query.include_forwarded
-                    and msg["forwarded_from"]
-                    and any(a in query.agents for a in msg["forwarded_from"])
-                )
-            ):
+            if query.agents and msg["name"] not in query.agents:
                 continue
 
             # Skip if before cutoff time
@@ -90,7 +111,6 @@ class MemoryStorageProvider(StorageProvider):
                 model_name=msg["model"],
                 cost_info=cost_info,
                 response_time=msg["response_time"],
-                forwarded_from=msg["forwarded_from"] or [],
                 timestamp=msg["timestamp"],
                 provider_name=msg["provider_name"],
                 provider_response_id=msg["provider_response_id"],
@@ -116,11 +136,11 @@ class MemoryStorageProvider(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,
         finish_reason: str | None = None,
+        parent_id: str | None = None,
     ) -> None:
         """Store message in memory."""
         if next((i for i in self.messages if i["message_id"] == message_id), None):
@@ -130,13 +150,13 @@ class MemoryStorageProvider(StorageProvider):
         self.messages.append({
             "conversation_id": conversation_id,
             "message_id": message_id,
+            "parent_id": parent_id,
             "content": content,
             "role": role,
             "name": name,
             "cost_info": cost_info.token_usage if cost_info else None,
             "model": model,
             "response_time": response_time,
-            "forwarded_from": forwarded_from,
             "provider_name": provider_name,
             "provider_response_id": provider_response_id,
             "messages": messages,
@@ -183,6 +203,111 @@ class MemoryStorageProvider(StorageProvider):
                 return conv.get("title")
         return None
 
+    async def get_conversation_messages(
+        self,
+        conversation_id: str,
+        *,
+        include_ancestors: bool = False,
+    ) -> list[ChatMessage[str]]:
+        """Get all messages for a conversation."""
+        messages: list[ChatMessage[str]] = []
+        for msg in self.messages:
+            if msg.get("conversation_id") != conversation_id:
+                continue
+            messages.append(_dict_to_chat_message(msg))
+
+        # Sort by timestamp
+        messages.sort(key=lambda m: m.timestamp or get_now())
+
+        if not include_ancestors or not messages:
+            return messages
+
+        # Get ancestor chain if first message has parent_id
+        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."""
+        for msg in self.messages:
+            if msg.get("message_id") == message_id:
+                return _dict_to_chat_message(msg)
+        return None
+
+    async def get_message_ancestry(self, message_id: str) -> list[ChatMessage[str]]:
+        """Get the ancestry chain of a message."""
+        ancestors: list[ChatMessage[str]] = []
+        current_id: str | None = message_id
+
+        while current_id:
+            msg = await self.get_message(current_id)
+            if not msg:
+                break
+            ancestors.append(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."""
+        # Find source conversation
+        source_conv = next(
+            (c for c in self.conversations if c["id"] == source_conversation_id), 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 message exists in source conversation
+            msg_exists = any(
+                m.get("message_id") == fork_from_message_id
+                and m["conversation_id"] == source_conversation_id
+                for m in self.messages
+            )
+            if not msg_exists:
+                err = f"Message {fork_from_message_id} not found in conversation"
+                raise ValueError(err)
+            fork_point_id = fork_from_message_id
+        else:
+            # Find last message in source conversation
+            conv_messages = [
+                m for m in self.messages if m["conversation_id"] == source_conversation_id
+            ]
+            if conv_messages:
+                conv_messages.sort(key=lambda m: m.get("timestamp") or get_now())
+                fork_point_id = conv_messages[-1].get("message_id")
+
+        # Create new conversation
+        agent_name = new_agent_name or source_conv["agent_name"]
+        title = (
+            f"{source_conv.get('title') or 'Conversation'} (fork)"
+            if source_conv.get("title")
+            else None
+        )
+        self.conversations.append({
+            "id": new_conversation_id,
+            "agent_name": agent_name,
+            "title": title,
+            "start_time": get_now(),
+        })
+
+        return fork_point_id
+
     async def log_command(
         self,
         *,
@@ -261,7 +386,6 @@ class MemoryStorageProvider(StorageProvider):
                     model_name=msg["model"],
                     cost_info=cost_info,
                     response_time=msg["response_time"],
-                    forwarded_from=msg["forwarded_from"],
                     timestamp=msg["timestamp"],
                 )
                 conv_messages.append(chat_msg)
@@ -394,3 +518,60 @@ class MemoryStorageProvider(StorageProvider):
             msg_count = len(self.messages)
 
         return conv_count, msg_count
+
+    async def delete_conversation_messages(
+        self,
+        conversation_id: str,
+    ) -> int:
+        """Delete all messages for a conversation."""
+        original_count = len(self.messages)
+        self.messages = [m for m in self.messages if m["conversation_id"] != conversation_id]
+        return original_count - len(self.messages)
+
+    # Project methods
+
+    async def save_project(self, project: ProjectData) -> None:
+        """Save or update a project."""
+        self.projects[project.project_id] = project
+
+    async def get_project(self, project_id: str) -> ProjectData | None:
+        """Get a project by ID."""
+        return self.projects.get(project_id)
+
+    async def get_project_by_worktree(self, worktree: str) -> ProjectData | None:
+        """Get a project by worktree path."""
+        for project in self.projects.values():
+            if project.worktree == worktree:
+                return project
+        return None
+
+    async def get_project_by_name(self, name: str) -> ProjectData | None:
+        """Get a project by friendly name."""
+        for project in self.projects.values():
+            if project.name == name:
+                return project
+        return None
+
+    async def list_projects(self, limit: int | None = None) -> list[ProjectData]:
+        """List all projects, ordered by last_active descending."""
+        projects = sorted(
+            self.projects.values(),
+            key=lambda p: p.last_active,
+            reverse=True,
+        )
+        if limit is not None:
+            projects = projects[:limit]
+        return list(projects)
+
+    async def delete_project(self, project_id: str) -> bool:
+        """Delete a project."""
+        if project_id in self.projects:
+            del self.projects[project_id]
+            return True
+        return False
+
+    async def touch_project(self, project_id: str) -> None:
+        """Update project's last_active timestamp."""
+        if project_id in self.projects:
+            project = self.projects[project_id]
+            self.projects[project_id] = project.touch()

agentpool_storage/models.py

@@ -36,6 +36,9 @@ class MessageData(TypedDict):
     timestamp: str
     """When the message was sent (ISO format)"""
 
+    parent_id: str | None
+    """ID of the parent message for tree-structured conversations."""
+
     model: str | None
     """Name of the model that generated this message"""