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

agentpool/storage/manager.py

@@ -9,6 +9,7 @@ from anyenv import method_spawner
 from pydantic_ai import Agent
 
 from agentpool.log import get_logger
+from agentpool.messaging import ChatMessage
 from agentpool.storage.serialization import serialize_messages
 from agentpool.utils.tasks import TaskManager
 from agentpool_config.storage import (
@@ -25,7 +26,7 @@ if TYPE_CHECKING:
     from types import TracebackType
 
     from agentpool.common_types import JsonValue
-    from agentpool.messaging import ChatMessage
+    from agentpool.sessions.models import ProjectData
     from agentpool_config.session import SessionQuery
     from agentpool_config.storage import BaseStorageProviderConfig, StorageConfig
     from agentpool_storage.base import StorageProvider
@@ -209,6 +210,7 @@ class StorageManager:
                     content=str(message.content),
                     role=message.role,
                     name=message.name,
+                    parent_id=message.parent_id,
                     cost_info=message.cost_info,
                     model=message.model_name,
                     response_time=message.response_time,
@@ -363,6 +365,87 @@ class StorageManager:
         provider = self.get_history_provider()
         return await provider.get_conversation_title(conversation_id)
 
+    @method_spawner
+    async def delete_conversation_messages(
+        self,
+        conversation_id: str,
+    ) -> int:
+        """Delete all messages for a conversation in all providers.
+
+        Used for compaction - removes existing messages so they can be
+        replaced with compacted versions.
+
+        Args:
+            conversation_id: ID of the conversation to clear
+
+        Returns:
+            Total number of messages deleted across all providers
+        """
+        total_deleted = 0
+        for provider in self.providers:
+            try:
+                deleted = await provider.delete_conversation_messages(conversation_id)
+                total_deleted += deleted
+            except NotImplementedError:
+                # Provider doesn't support deletion (e.g., write-only log)
+                pass
+            except Exception:
+                logger.exception(
+                    "Error deleting messages from provider",
+                    provider=provider.__class__.__name__,
+                    conversation_id=conversation_id,
+                )
+        return total_deleted
+
+    @method_spawner
+    async def replace_conversation_messages(
+        self,
+        conversation_id: str,
+        messages: Sequence[ChatMessage[Any]],
+    ) -> tuple[int, int]:
+        """Replace all messages for a conversation with new ones.
+
+        Deletes existing messages and logs new ones. Used for compaction
+        where the full history is replaced with a compacted version.
+
+        Args:
+            conversation_id: ID of the conversation
+            messages: New messages to store
+
+        Returns:
+            Tuple of (deleted_count, added_count)
+        """
+        # First delete existing messages
+        deleted = await self.delete_conversation_messages(conversation_id)
+
+        # Then log new messages
+        added = 0
+        for message in messages:
+            # Ensure conversation_id is set on the message
+            msg_to_log: ChatMessage[Any] = message
+            if not message.conversation_id:
+                msg_to_log = ChatMessage(
+                    content=message.content,
+                    role=message.role,
+                    name=message.name,
+                    conversation_id=conversation_id,
+                    message_id=message.message_id,
+                    parent_id=message.parent_id,
+                    model_name=message.model_name,
+                    cost_info=message.cost_info,
+                    response_time=message.response_time,
+                    forwarded_from=message.forwarded_from,
+                    timestamp=message.timestamp,
+                    provider_name=message.provider_name,
+                    provider_response_id=message.provider_response_id,
+                    messages=message.messages,
+                    finish_reason=message.finish_reason,
+                )
+            await self.log_message(msg_to_log)
+            added += 1
+
+        return deleted, added
+
     async def generate_conversation_title(
         self,
         conversation_id: str,
@@ -417,3 +500,136 @@ class StorageManager:
             return None
         else:
             return title
+
+    # Project methods
+
+    def get_project_provider(self) -> StorageProvider:
+        """Get provider capable of storing projects.
+
+        Returns:
+            First provider that supports project storage.
+
+        Raises:
+            RuntimeError: If no capable provider found.
+        """
+        for provider in self.providers:
+            # TextLogProvider doesn't support projects, others do
+            if hasattr(provider, "save_project") and not getattr(provider, "write_only", False):
+                return provider
+        msg = "No provider found that supports project storage"
+        raise RuntimeError(msg)
+
+    @method_spawner
+    async def save_project(self, project: ProjectData) -> None:
+        """Save or update a project in all capable providers.
+
+        Args:
+            project: Project data to persist
+        """
+        for provider in self.providers:
+            try:
+                await provider.save_project(project)
+            except NotImplementedError:
+                pass
+            except Exception:
+                logger.exception(
+                    "Error saving project",
+                    provider=provider.__class__.__name__,
+                    project_id=project.project_id,
+                )
+
+    @method_spawner
+    async def get_project(self, project_id: str) -> ProjectData | None:
+        """Get a project by ID.
+
+        Args:
+            project_id: Project identifier
+
+        Returns:
+            Project data if found, None otherwise
+        """
+        provider = self.get_project_provider()
+        return await provider.get_project(project_id)
+
+    @method_spawner
+    async def get_project_by_worktree(self, worktree: str) -> ProjectData | None:
+        """Get a project by worktree path.
+
+        Args:
+            worktree: Absolute path to the project worktree
+
+        Returns:
+            Project data if found, None otherwise
+        """
+        provider = self.get_project_provider()
+        return await provider.get_project_by_worktree(worktree)
+
+    @method_spawner
+    async def get_project_by_name(self, name: str) -> ProjectData | None:
+        """Get a project by friendly name.
+
+        Args:
+            name: Project name
+
+        Returns:
+            Project data if found, None otherwise
+        """
+        provider = self.get_project_provider()
+        return await provider.get_project_by_name(name)
+
+    @method_spawner
+    async def list_projects(self, limit: int | None = None) -> list[ProjectData]:
+        """List all projects, ordered by last_active descending.
+
+        Args:
+            limit: Maximum number of projects to return
+
+        Returns:
+            List of project data objects
+        """
+        provider = self.get_project_provider()
+        return await provider.list_projects(limit=limit)
+
+    @method_spawner
+    async def delete_project(self, project_id: str) -> bool:
+        """Delete a project from all providers.
+
+        Args:
+            project_id: Project identifier
+
+        Returns:
+            True if project was deleted from at least one provider
+        """
+        deleted = False
+        for provider in self.providers:
+            try:
+                if await provider.delete_project(project_id):
+                    deleted = True
+            except NotImplementedError:
+                pass
+            except Exception:
+                logger.exception(
+                    "Error deleting project",
+                    provider=provider.__class__.__name__,
+                    project_id=project_id,
+                )
+        return deleted
+
+    @method_spawner
+    async def touch_project(self, project_id: str) -> None:
+        """Update project's last_active timestamp in all providers.
+
+        Args:
+            project_id: Project identifier
+        """
+        for provider in self.providers:
+            try:
+                await provider.touch_project(project_id)
+            except NotImplementedError:
+                pass
+            except Exception:
+                logger.exception(
+                    "Error touching project",
+                    provider=provider.__class__.__name__,
+                    project_id=project_id,
+                )