dooers-workers 0.2.0__py3-none-any.whl

dooers/persistence/sqlite.py

@@ -0,0 +1,433 @@
+import json
+from datetime import UTC, datetime
+from typing import Any
+
+import aiosqlite
+
+from dooers.protocol.models import DocumentPart, ImagePart, Run, TextPart, Thread, ThreadEvent
+
+
+class SqlitePersistence:
+    def __init__(self, database_url: str, table_prefix: str = "worker_"):
+        self._database_url = database_url.replace("sqlite:///", "")
+        self._prefix = table_prefix
+        self._conn: aiosqlite.Connection | None = None
+
+    async def connect(self) -> None:
+        self._conn = await aiosqlite.connect(self._database_url)
+        self._conn.row_factory = aiosqlite.Row
+
+    async def disconnect(self) -> None:
+        if self._conn:
+            await self._conn.close()
+
+    async def migrate(self) -> None:
+        if not self._conn:
+            raise RuntimeError("Not connected")
+
+        threads_table = f"{self._prefix}threads"
+        events_table = f"{self._prefix}events"
+        runs_table = f"{self._prefix}runs"
+
+        await self._conn.executescript(f"""
+            CREATE TABLE IF NOT EXISTS {threads_table} (
+                id TEXT PRIMARY KEY,
+                worker_id TEXT NOT NULL,
+                user_id TEXT,
+                title TEXT,
+                created_at TEXT NOT NULL,
+                updated_at TEXT NOT NULL,
+                last_event_at TEXT NOT NULL
+            );
+
+            CREATE TABLE IF NOT EXISTS {events_table} (
+                id TEXT PRIMARY KEY,
+                thread_id TEXT NOT NULL,
+                run_id TEXT,
+                type TEXT NOT NULL,
+                actor TEXT NOT NULL,
+                user_id TEXT,
+                user_name TEXT,
+                user_email TEXT,
+                content TEXT,
+                data TEXT,
+                created_at TEXT NOT NULL,
+                FOREIGN KEY (thread_id) REFERENCES {threads_table}(id)
+            );
+
+            CREATE TABLE IF NOT EXISTS {runs_table} (
+                id TEXT PRIMARY KEY,
+                thread_id TEXT NOT NULL,
+                agent_id TEXT,
+                status TEXT NOT NULL,
+                started_at TEXT NOT NULL,
+                ended_at TEXT,
+                error TEXT,
+                FOREIGN KEY (thread_id) REFERENCES {threads_table}(id)
+            );
+
+            CREATE INDEX IF NOT EXISTS idx_{self._prefix}threads_worker_id
+                ON {threads_table}(worker_id);
+            CREATE INDEX IF NOT EXISTS idx_{self._prefix}threads_user_id
+                ON {threads_table}(user_id);
+            CREATE INDEX IF NOT EXISTS idx_{self._prefix}events_thread_id
+                ON {events_table}(thread_id);
+            CREATE INDEX IF NOT EXISTS idx_{self._prefix}events_user_id
+                ON {events_table}(user_id);
+            CREATE INDEX IF NOT EXISTS idx_{self._prefix}runs_thread_id
+                ON {runs_table}(thread_id);
+
+            CREATE TABLE IF NOT EXISTS {self._prefix}settings (
+                worker_id TEXT PRIMARY KEY,
+                values TEXT NOT NULL DEFAULT '{{}}',
+                created_at TEXT DEFAULT (datetime('now')),
+                updated_at TEXT DEFAULT (datetime('now'))
+            );
+            CREATE INDEX IF NOT EXISTS idx_{self._prefix}settings_worker
+                ON {self._prefix}settings(worker_id);
+        """)
+        await self._conn.commit()
+
+    async def create_thread(self, thread: Thread) -> None:
+        if not self._conn:
+            raise RuntimeError("Not connected")
+
+        table = f"{self._prefix}threads"
+        await self._conn.execute(
+            f"""
+            INSERT INTO {table} (id, worker_id, user_id, title, created_at, updated_at, last_event_at)
+            VALUES (?, ?, ?, ?, ?, ?, ?)
+            """,
+            (
+                thread.id,
+                thread.worker_id,
+                thread.user_id,
+                thread.title,
+                thread.created_at.isoformat(),
+                thread.updated_at.isoformat(),
+                thread.last_event_at.isoformat(),
+            ),
+        )
+        await self._conn.commit()
+
+    async def get_thread(self, thread_id: str) -> Thread | None:
+        if not self._conn:
+            raise RuntimeError("Not connected")
+
+        table = f"{self._prefix}threads"
+        cursor = await self._conn.execute(
+            f"SELECT * FROM {table} WHERE id = ?",
+            (thread_id,),
+        )
+        row = await cursor.fetchone()
+
+        if not row:
+            return None
+
+        return Thread(
+            id=row["id"],
+            worker_id=row["worker_id"],
+            user_id=row["user_id"],
+            title=row["title"],
+            created_at=datetime.fromisoformat(row["created_at"]),
+            updated_at=datetime.fromisoformat(row["updated_at"]),
+            last_event_at=datetime.fromisoformat(row["last_event_at"]),
+        )
+
+    async def update_thread(self, thread: Thread) -> None:
+        if not self._conn:
+            raise RuntimeError("Not connected")
+
+        table = f"{self._prefix}threads"
+        await self._conn.execute(
+            f"""
+            UPDATE {table}
+            SET user_id = ?, title = ?, updated_at = ?, last_event_at = ?
+            WHERE id = ?
+            """,
+            (
+                thread.user_id,
+                thread.title,
+                thread.updated_at.isoformat(),
+                thread.last_event_at.isoformat(),
+                thread.id,
+            ),
+        )
+        await self._conn.commit()
+
+    async def list_threads(
+        self,
+        worker_id: str,
+        user_id: str | None,
+        cursor: str | None,
+        limit: int,
+    ) -> list[Thread]:
+        if not self._conn:
+            raise RuntimeError("Not connected")
+
+        table = f"{self._prefix}threads"
+
+        if user_id:
+            query = f"""
+                SELECT * FROM {table}
+                WHERE worker_id = ? AND user_id = ?
+                ORDER BY last_event_at DESC
+                LIMIT ?
+            """
+            cursor_result = await self._conn.execute(query, (worker_id, user_id, limit))
+        else:
+            query = f"""
+                SELECT * FROM {table}
+                WHERE worker_id = ?
+                ORDER BY last_event_at DESC
+                LIMIT ?
+            """
+            cursor_result = await self._conn.execute(query, (worker_id, limit))
+
+        rows = await cursor_result.fetchall()
+
+        return [
+            Thread(
+                id=row["id"],
+                worker_id=row["worker_id"],
+                user_id=row["user_id"],
+                title=row["title"],
+                created_at=datetime.fromisoformat(row["created_at"]),
+                updated_at=datetime.fromisoformat(row["updated_at"]),
+                last_event_at=datetime.fromisoformat(row["last_event_at"]),
+            )
+            for row in rows
+        ]
+
+    async def create_event(self, event: ThreadEvent) -> None:
+        if not self._conn:
+            raise RuntimeError("Not connected")
+
+        table = f"{self._prefix}events"
+        content_json = None
+        if event.content:
+            content_json = json.dumps([self._serialize_content_part(p) for p in event.content])
+
+        data_json = json.dumps(event.data) if event.data else None
+
+        await self._conn.execute(
+            f"""
+            INSERT INTO {table} (id, thread_id, run_id, type, actor, user_id, user_name, user_email, content, data, created_at)
+            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
+            """,
+            (
+                event.id,
+                event.thread_id,
+                event.run_id,
+                event.type,
+                event.actor,
+                event.user_id,
+                event.user_name,
+                event.user_email,
+                content_json,
+                data_json,
+                event.created_at.isoformat(),
+            ),
+        )
+        await self._conn.commit()
+
+    async def get_events(
+        self,
+        thread_id: str,
+        after_event_id: str | None,
+        limit: int,
+    ) -> list[ThreadEvent]:
+        if not self._conn:
+            raise RuntimeError("Not connected")
+
+        table = f"{self._prefix}events"
+
+        if after_event_id:
+            cursor_result = await self._conn.execute(
+                f"SELECT created_at FROM {table} WHERE id = ?",
+                (after_event_id,),
+            )
+            ref_row = await cursor_result.fetchone()
+            if ref_row:
+                query = f"""
+                    SELECT * FROM {table}
+                    WHERE thread_id = ? AND created_at > ?
+                    ORDER BY created_at ASC
+                    LIMIT ?
+                """
+                cursor_result = await self._conn.execute(
+                    query,
+                    (thread_id, ref_row["created_at"], limit),
+                )
+            else:
+                query = f"""
+                    SELECT * FROM {table}
+                    WHERE thread_id = ?
+                    ORDER BY created_at ASC
+                    LIMIT ?
+                """
+                cursor_result = await self._conn.execute(query, (thread_id, limit))
+        else:
+            query = f"""
+                SELECT * FROM {table}
+                WHERE thread_id = ?
+                ORDER BY created_at ASC
+                LIMIT ?
+            """
+            cursor_result = await self._conn.execute(query, (thread_id, limit))
+
+        rows = await cursor_result.fetchall()
+
+        events = []
+        for row in rows:
+            content = None
+            if row["content"]:
+                content_data = json.loads(row["content"])
+                content = [self._deserialize_content_part(p) for p in content_data]
+
+            data = json.loads(row["data"]) if row["data"] else None
+
+            events.append(
+                ThreadEvent(
+                    id=row["id"],
+                    thread_id=row["thread_id"],
+                    run_id=row["run_id"],
+                    type=row["type"],
+                    actor=row["actor"],
+                    user_id=row["user_id"],
+                    user_name=row["user_name"],
+                    user_email=row["user_email"],
+                    content=content,
+                    data=data,
+                    created_at=datetime.fromisoformat(row["created_at"]),
+                )
+            )
+
+        return events
+
+    async def create_run(self, run: Run) -> None:
+        if not self._conn:
+            raise RuntimeError("Not connected")
+
+        table = f"{self._prefix}runs"
+        await self._conn.execute(
+            f"""
+            INSERT INTO {table} (id, thread_id, agent_id, status, started_at, ended_at, error)
+            VALUES (?, ?, ?, ?, ?, ?, ?)
+            """,
+            (
+                run.id,
+                run.thread_id,
+                run.agent_id,
+                run.status,
+                run.started_at.isoformat(),
+                run.ended_at.isoformat() if run.ended_at else None,
+                run.error,
+            ),
+        )
+        await self._conn.commit()
+
+    async def update_run(self, run: Run) -> None:
+        if not self._conn:
+            raise RuntimeError("Not connected")
+
+        table = f"{self._prefix}runs"
+        await self._conn.execute(
+            f"""
+            UPDATE {table}
+            SET agent_id = ?, status = ?, ended_at = ?, error = ?
+            WHERE id = ?
+            """,
+            (
+                run.agent_id,
+                run.status,
+                run.ended_at.isoformat() if run.ended_at else None,
+                run.error,
+                run.id,
+            ),
+        )
+        await self._conn.commit()
+
+    def _serialize_content_part(self, part) -> dict:
+        if hasattr(part, "model_dump"):
+            return part.model_dump()
+        return dict(part)
+
+    def _deserialize_content_part(self, data: dict):
+        part_type = data.get("type")
+        if part_type == "text":
+            return TextPart(**data)
+        elif part_type == "image":
+            return ImagePart(**data)
+        elif part_type == "document":
+            return DocumentPart(**data)
+        return data
+
+    async def get_settings(self, worker_id: str) -> dict[str, Any]:
+        """Get all stored values for a worker. Returns empty dict if none."""
+        if not self._conn:
+            raise RuntimeError("Not connected")
+
+        table = f"{self._prefix}settings"
+        cursor = await self._conn.execute(
+            f"SELECT values FROM {table} WHERE worker_id = ?",
+            (worker_id,),
+        )
+        row = await cursor.fetchone()
+
+        if not row:
+            return {}
+
+        return json.loads(row["values"])
+
+    async def update_setting(self, worker_id: str, field_id: str, value: Any) -> datetime:
+        """Update a single field value. Returns updated_at timestamp."""
+        if not self._conn:
+            raise RuntimeError("Not connected")
+
+        table = f"{self._prefix}settings"
+        now = datetime.now(UTC)
+        now_str = now.isoformat()
+
+        # Get existing values
+        current_values = await self.get_settings(worker_id)
+        current_values[field_id] = value
+        values_json = json.dumps(current_values)
+
+        # Upsert
+        await self._conn.execute(
+            f"""
+            INSERT INTO {table} (worker_id, values, created_at, updated_at)
+            VALUES (?, ?, ?, ?)
+            ON CONFLICT(worker_id) DO UPDATE SET
+                values = excluded.values,
+                updated_at = excluded.updated_at
+            """,
+            (worker_id, values_json, now_str, now_str),
+        )
+        await self._conn.commit()
+        return now
+
+    async def set_settings(self, worker_id: str, values: dict[str, Any]) -> datetime:
+        """Replace all settings values. Returns updated_at timestamp."""
+        if not self._conn:
+            raise RuntimeError("Not connected")
+
+        table = f"{self._prefix}settings"
+        now = datetime.now(UTC)
+        now_str = now.isoformat()
+        values_json = json.dumps(values)
+
+        # Upsert
+        await self._conn.execute(
+            f"""
+            INSERT INTO {table} (worker_id, values, created_at, updated_at)
+            VALUES (?, ?, ?, ?)
+            ON CONFLICT(worker_id) DO UPDATE SET
+                values = excluded.values,
+                updated_at = excluded.updated_at
+            """,
+            (worker_id, values_json, now_str, now_str),
+        )
+        await self._conn.commit()
+        return now

dooers/protocol/__init__.py

@@ -0,0 +1,108 @@
+from dooers.protocol.frames import (
+    # Payload types
+    AnalyticsSubscribePayload,
+    AnalyticsUnsubscribePayload,
+    # Analytics C2S frames
+    C2S_AnalyticsSubscribe,
+    C2S_AnalyticsUnsubscribe,
+    # Existing C2S frames
+    C2S_Connect,
+    C2S_EventCreate,
+    C2S_Feedback,
+    # Settings C2S frames
+    C2S_SettingsPatch,
+    C2S_SettingsSubscribe,
+    C2S_SettingsUnsubscribe,
+    C2S_ThreadList,
+    C2S_ThreadSubscribe,
+    C2S_ThreadUnsubscribe,
+    FeedbackAckPayload,
+    FeedbackPayload,
+    # Existing S2C frames
+    S2C_Ack,
+    # Analytics S2C frames
+    S2C_AnalyticsEvent,
+    S2C_EventAppend,
+    S2C_FeedbackAck,
+    S2C_RunUpsert,
+    # Settings S2C frames
+    S2C_SettingsPatch,
+    S2C_SettingsSnapshot,
+    S2C_ThreadListResult,
+    S2C_ThreadSnapshot,
+    S2C_ThreadUpsert,
+    SettingsPatchBroadcastPayload,
+    SettingsPatchPayload,
+    SettingsSnapshotPayload,
+    SettingsSubscribePayload,
+    SettingsUnsubscribePayload,
+    WSFrame,
+)
+from dooers.protocol.models import (
+    Actor,
+    ContentPart,
+    DocumentPart,
+    EventType,
+    ImagePart,
+    Run,
+    RunStatus,
+    TextPart,
+    Thread,
+    ThreadEvent,
+)
+from dooers.protocol.parser import parse_frame, serialize_frame
+
+__all__ = [
+    # Models
+    "ContentPart",
+    "TextPart",
+    "ImagePart",
+    "DocumentPart",
+    "Thread",
+    "ThreadEvent",
+    "Run",
+    "RunStatus",
+    "Actor",
+    "EventType",
+    "WSFrame",
+    # Existing C2S frames
+    "C2S_Connect",
+    "C2S_ThreadList",
+    "C2S_ThreadSubscribe",
+    "C2S_ThreadUnsubscribe",
+    "C2S_EventCreate",
+    # Analytics C2S frames
+    "C2S_AnalyticsSubscribe",
+    "C2S_AnalyticsUnsubscribe",
+    "C2S_Feedback",
+    # Settings C2S frames
+    "C2S_SettingsSubscribe",
+    "C2S_SettingsUnsubscribe",
+    "C2S_SettingsPatch",
+    # Existing S2C frames
+    "S2C_Ack",
+    "S2C_ThreadListResult",
+    "S2C_ThreadSnapshot",
+    "S2C_EventAppend",
+    "S2C_ThreadUpsert",
+    "S2C_RunUpsert",
+    # Analytics S2C frames
+    "S2C_AnalyticsEvent",
+    "S2C_FeedbackAck",
+    # Settings S2C frames
+    "S2C_SettingsSnapshot",
+    "S2C_SettingsPatch",
+    # Payload types
+    "AnalyticsSubscribePayload",
+    "AnalyticsUnsubscribePayload",
+    "FeedbackPayload",
+    "FeedbackAckPayload",
+    "SettingsSubscribePayload",
+    "SettingsUnsubscribePayload",
+    "SettingsPatchPayload",
+    "SettingsSnapshotPayload",
+    "SettingsPatchBroadcastPayload",
+    # Parser
+    "parse_frame",
+    "serialize_frame",
+]