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

agentpool_storage/file_provider.py

@@ -22,6 +22,7 @@ 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
 
@@ -46,6 +47,7 @@ class MessageData(TypedDict):
     provider_response_id: str | None
     messages: str | None
     finish_reason: FinishReason | None
+    parent_id: str | None
 
 
 class ConversationData(TypedDict):
@@ -68,12 +70,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 +115,7 @@ class FileProvider(StorageProvider):
             "messages": [],
             "conversations": [],
             "commands": [],
+            "projects": [],
         }
         self._load()
 
@@ -205,6 +222,7 @@ class FileProvider(StorageProvider):
         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({
@@ -227,6 +245,7 @@ class FileProvider(StorageProvider):
             "provider_response_id": provider_response_id,
             "messages": messages,
             "finish_reason": finish_reason,
+            "parent_id": parent_id,
         })
         self._save()
 
@@ -329,6 +348,7 @@ class FileProvider(StorageProvider):
                 "messages": [],
                 "conversations": [],
                 "commands": [],
+                "projects": [],
             }
             self._save()
             return conv_count, msg_count
@@ -376,3 +396,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,6 +16,7 @@ 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
@@ -31,12 +32,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."""
@@ -121,6 +124,7 @@ class MemoryStorageProvider(StorageProvider):
         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):
@@ -394,3 +398,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"""