flock-core 0.4.542__py3-none-any.whl → 0.5.0__py3-none-any.whl

flock/webapp/app/services/sharing_store.py

@@ -1,597 +0,0 @@
-"""Shared link and feedback storage implementations supporting SQLite and Azure Table Storage."""
-
-import logging
-import sqlite3
-from abc import ABC, abstractmethod
-from pathlib import Path
-from typing import Any
-
-import aiosqlite
-
-from flock.webapp.app.services.sharing_models import (
-    FeedbackRecord,
-    SharedLinkConfig,
-)
-
-# Azure Table Storage imports - will be conditionally imported
-try:
-    from azure.core.exceptions import ResourceExistsError, ResourceNotFoundError
-    from azure.data.tables.aio import TableServiceClient
-    AZURE_AVAILABLE = True
-except ImportError:
-    AZURE_AVAILABLE = False
-    TableServiceClient = None
-    ResourceNotFoundError = None
-    ResourceExistsError = None
-
-# Get a logger instance
-logger = logging.getLogger(__name__)
-
-class SharedLinkStoreInterface(ABC):
-    """Interface for storing and retrieving shared link configurations."""
-
-    @abstractmethod
-    async def initialize(self) -> None:
-        """Initialize the store (e.g., create tables)."""
-        pass
-
-    @abstractmethod
-    async def save_config(self, config: SharedLinkConfig) -> SharedLinkConfig:
-        """Saves a shared link configuration."""
-        pass
-
-    @abstractmethod
-    async def get_config(self, share_id: str) -> SharedLinkConfig | None:
-        """Retrieves a shared link configuration by its ID."""
-        pass
-
-    @abstractmethod
-    async def delete_config(self, share_id: str) -> bool:
-        """Deletes a shared link configuration by its ID. Returns True if deleted, False otherwise."""
-        pass
-
-    # Feedback
-    @abstractmethod
-    async def save_feedback(self, record: FeedbackRecord):
-        """Persist a feedback record."""
-        pass
-
-    @abstractmethod
-    async def get_feedback(self, id: str) -> FeedbackRecord | None:
-        """Get a single feedback record."""
-        pass
-
-    @abstractmethod
-    async def get_all_feedback_records_for_agent(self, agent_name: str) -> list[FeedbackRecord]:
-        """Get all feedback records for a given agent."""
-        pass
-
-class SQLiteSharedLinkStore(SharedLinkStoreInterface):
-    """SQLite implementation for storing and retrieving shared link configurations."""
-
-    def __init__(self, db_path: str):
-        """Initialize SQLite store with database path."""
-        self.db_path = Path(db_path)
-        self.db_path.parent.mkdir(parents=True, exist_ok=True) # Ensure directory exists
-        logger.info(f"SQLiteSharedLinkStore initialized with db_path: {self.db_path}")
-
-    async def initialize(self) -> None:
-        """Initializes the database and creates/updates the table if it doesn't exist."""
-        try:
-            async with aiosqlite.connect(self.db_path) as db:
-                # Ensure the table exists with the base schema first
-                await db.execute(
-                    """
-                    CREATE TABLE IF NOT EXISTS shared_links (
-                        share_id TEXT PRIMARY KEY,
-                        agent_name TEXT NOT NULL,
-                        flock_definition TEXT NOT NULL,
-                        created_at TEXT NOT NULL
-                        /* New columns will be added below if they don't exist */
-                    )
-                    """
-                )
-
-                # Add new columns individually, ignoring errors if they already exist
-                new_columns = [
-                    ("share_type", "TEXT DEFAULT 'agent_run' NOT NULL"),
-                    ("chat_message_key", "TEXT"),
-                    ("chat_history_key", "TEXT"),
-                    ("chat_response_key", "TEXT")
-                ]
-
-                for column_name, column_type in new_columns:
-                    try:
-                        await db.execute(f"ALTER TABLE shared_links ADD COLUMN {column_name} {column_type}")
-                        logger.info(f"Added column '{column_name}' to shared_links table.")
-                    except sqlite3.OperationalError as e:
-                        if "duplicate column name" in str(e).lower():
-                            logger.debug(f"Column '{column_name}' already exists in shared_links table.")
-                        else:
-                            raise # Re-raise if it's a different operational error
-
-                # Feedback table
-                await db.execute(
-                    """
-                    CREATE TABLE IF NOT EXISTS feedback (
-                        feedback_id TEXT PRIMARY KEY,
-                        share_id TEXT,
-                        context_type TEXT NOT NULL,
-                        reason TEXT NOT NULL,
-                        expected_response TEXT,
-                        actual_response TEXT,
-                        flock_name TEXT,
-                        agent_name TEXT,
-                        flock_definition TEXT,
-                        created_at TEXT NOT NULL,
-                        FOREIGN KEY(share_id) REFERENCES shared_links(share_id)
-                    )
-                    """
-                )
-
-                await db.commit()
-            logger.info(f"Database initialized and shared_links table schema ensured at {self.db_path}")
-        except sqlite3.Error as e:
-            logger.error(f"SQLite error during initialization: {e}", exc_info=True)
-            raise
-
-    async def save_config(self, config: SharedLinkConfig) -> SharedLinkConfig:
-        """Saves a shared link configuration to the SQLite database."""
-        try:
-            async with aiosqlite.connect(self.db_path) as db:
-                await db.execute(
-                    """INSERT INTO shared_links (
-                        share_id, agent_name, created_at, flock_definition, 
-                        share_type, chat_message_key, chat_history_key, chat_response_key
-                    ) VALUES (?, ?, ?, ?, ?, ?, ?, ?)""",
-                    (
-                        config.share_id,
-                        config.agent_name,
-                        config.created_at.isoformat(),
-                        config.flock_definition,
-                        config.share_type,
-                        config.chat_message_key,
-                        config.chat_history_key,
-                        config.chat_response_key,
-                    ),
-                )
-                await db.commit()
-            logger.info(f"Saved shared link config for ID: {config.share_id} with type: {config.share_type}")
-            return config
-        except sqlite3.Error as e:
-            logger.error(f"SQLite error saving config for ID {config.share_id}: {e}", exc_info=True)
-            raise
-
-    async def get_config(self, share_id: str) -> SharedLinkConfig | None:
-        """Retrieves a shared link configuration from SQLite by its ID."""
-        try:
-            async with aiosqlite.connect(self.db_path) as db:
-                async with db.execute(
-                    """SELECT 
-                        share_id, agent_name, created_at, flock_definition, 
-                        share_type, chat_message_key, chat_history_key, chat_response_key 
-                    FROM shared_links WHERE share_id = ?""",
-                    (share_id,)
-                ) as cursor:
-                    row = await cursor.fetchone()
-            if row:
-                logger.debug(f"Retrieved shared link config for ID: {share_id}")
-                return SharedLinkConfig(
-                    share_id=row[0],
-                    agent_name=row[1],
-                    created_at=row[2], # SQLite stores as TEXT, Pydantic will parse from ISO format
-                    flock_definition=row[3],
-                    share_type=row[4],
-                    chat_message_key=row[5],
-                    chat_history_key=row[6],
-                    chat_response_key=row[7],
-                )
-            logger.debug(f"No shared link config found for ID: {share_id}")
-            return None
-        except sqlite3.Error as e:
-            logger.error(f"SQLite error retrieving config for ID {share_id}: {e}", exc_info=True)
-            return None # Or raise, depending on desired error handling
-
-    async def delete_config(self, share_id: str) -> bool:
-        """Deletes a shared link configuration from SQLite by its ID."""
-        try:
-            async with aiosqlite.connect(self.db_path) as db:
-                result = await db.execute("DELETE FROM shared_links WHERE share_id = ?", (share_id,))
-                await db.commit()
-                deleted_count = result.rowcount
-            if deleted_count > 0:
-                logger.info(f"Deleted shared link config for ID: {share_id}")
-                return True
-            logger.info(f"Attempted to delete non-existent shared link config for ID: {share_id}")
-            return False
-        except sqlite3.Error as e:
-            logger.error(f"SQLite error deleting config for ID {share_id}: {e}", exc_info=True)
-            return False # Or raise
-
-    # ----------------------- Feedback methods -----------------------
-    async def get_feedback(self, id: str) -> FeedbackRecord | None:
-        """Retrieve a single feedback record from SQLite."""
-        try:
-            async with aiosqlite.connect(self.db_path) as db, db.execute(
-                """SELECT
-                    feedback_id, share_id, context_type, reason,
-                    expected_response, actual_response, flock_name, agent_name, flock_definition, created_at
-                FROM feedback WHERE feedback_id = ?""",
-                (id,)
-            ) as cursor:
-                row = await cursor.fetchone()
-
-            if row:
-                logger.debug(f"Retrieved feedback record for ID: {id}")
-                return FeedbackRecord(
-                    feedback_id=row[0],
-                    share_id=row[1],
-                    context_type=row[2],
-                    reason=row[3],
-                    expected_response=row[4],
-                    actual_response=row[5],
-                    flock_name=row[6],
-                    agent_name=row[7],
-                    flock_definition=row[8],
-                    created_at=row[9],  # SQLite stores as TEXT, Pydantic will parse from ISO format
-                )
-
-            logger.debug(f"No feedback record found for ID: {id}")
-            return None
-        except sqlite3.Error as e:
-            logger.error(f"SQLite error retrieving feedback for ID {id}: {e}", exc_info=True)
-            return None  # Or raise, depending on desired error handling
-
-    async def get_all_feedback_records_for_agent(self, agent_name: str) -> list[FeedbackRecord]:
-        """Retrieve all feedback records from SQLite."""
-        try:
-            async with aiosqlite.connect(self.db_path) as db, db.execute(
-                """SELECT
-                    feedback_id, share_id, context_type, reason,
-                    expected_response, actual_response, flock_name, agent_name, flock_definition, created_at
-                FROM feedback WHERE agent_name = ? ORDER BY created_at DESC""",
-                (agent_name,)
-            ) as cursor:
-                rows = await cursor.fetchall()
-
-            records = []
-            for row in rows:
-                records.append(FeedbackRecord(
-                    feedback_id=row[0],
-                    share_id=row[1],
-                    context_type=row[2],
-                    reason=row[3],
-                    expected_response=row[4],
-                    actual_response=row[5],
-                    flock_name=row[6],
-                    agent_name=row[7],
-                    flock_definition=row[8],
-                    created_at=row[9],  # SQLite stores as TEXT, Pydantic will parse from ISO format
-                ))
-
-            logger.debug(f"Retrieved {len(records)} feedback records")
-            return records
-        except sqlite3.Error as e:
-            logger.error(f"SQLite error retrieving all feedback records: {e}", exc_info=True)
-            return []  # Return empty list on error
-
-    async def save_feedback(self, record: FeedbackRecord) -> FeedbackRecord:
-        """Persist a feedback record to SQLite."""
-        try:
-            async with aiosqlite.connect(self.db_path) as db:
-                await db.execute(
-                    """INSERT INTO feedback (
-                        feedback_id, share_id, context_type, reason,
-                        expected_response, actual_response, flock_name, agent_name, flock_definition, created_at
-                    ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
-                    (
-                        record.feedback_id,
-                        record.share_id,
-                        record.context_type,
-                        record.reason,
-                        record.expected_response,
-                        record.actual_response,
-                        record.flock_name,
-                        record.agent_name,
-                        record.flock_definition,
-                        record.created_at.isoformat(),
-                    ),
-                )
-                await db.commit()
-            logger.info(f"Saved feedback {record.feedback_id} (share={record.share_id})")
-            return record
-        except sqlite3.Error as e:
-            logger.error(f"SQLite error saving feedback {record.feedback_id}: {e}", exc_info=True)
-            raise
-
-# ---------------------------------------------------------------------------
-# Azure Table + Blob implementation
-# ---------------------------------------------------------------------------
-
-try:
-    from azure.storage.blob.aio import BlobServiceClient
-    AZURE_BLOB_AVAILABLE = True
-except ImportError:  # blob SDK not installed
-    AZURE_BLOB_AVAILABLE = False
-    BlobServiceClient = None
-
-class AzureTableSharedLinkStore(SharedLinkStoreInterface):
-    """Store configs in Azure Table; store large flock YAML in Blob Storage."""
-
-    _TABLE_NAME        = "flocksharedlinks"
-    _FEEDBACK_TBL_NAME = "flockfeedback"
-    _CONTAINER_NAME    = "flocksharedlinkdefs"          # blobs live here
-    _PARTITION_KEY     = "shared_links"
-
-    def __init__(self, connection_string: str):
-        if not AZURE_AVAILABLE:
-            raise ImportError("pip install azure-data-tables")
-        if not AZURE_BLOB_AVAILABLE:
-            raise ImportError("pip install azure-storage-blob")
-
-        self.connection_string = connection_string
-        self.table_svc = TableServiceClient.from_connection_string(connection_string)
-        self.blob_svc  = BlobServiceClient.from_connection_string(connection_string)
-
-    # ------------------------------------------------------------------ init
-    async def initialize(self) -> None:
-        # 1. Azure Tables ----------------------------------------------------
-        try:
-            await self.table_svc.create_table(self._TABLE_NAME)
-            logger.info("Created Azure Table '%s'", self._TABLE_NAME)
-        except ResourceExistsError:
-            logger.debug("Azure Table '%s' already exists", self._TABLE_NAME)
-
-        try:
-            await self.table_svc.create_table(self._FEEDBACK_TBL_NAME)
-            logger.info("Created Azure Table '%s'", self._FEEDBACK_TBL_NAME)
-        except ResourceExistsError:
-            logger.debug("Azure Table '%s' already exists", self._FEEDBACK_TBL_NAME)
-
-        # 2. Blob container --------------------------------------------------
-        try:
-            await self.blob_svc.create_container(self._CONTAINER_NAME)
-            logger.info("Created Blob container '%s'", self._CONTAINER_NAME)
-        except ResourceExistsError:
-            logger.debug("Blob container '%s' already exists", self._CONTAINER_NAME)
-
-    # ------------------------------------------------------------- save_config
-    async def save_config(self, config: SharedLinkConfig) -> SharedLinkConfig:
-        """Upload YAML to Blob, then upsert table row containing the blob name."""
-        blob_name   = f"{config.share_id}.yaml"
-        blob_client = self.blob_svc.get_blob_client(self._CONTAINER_NAME, blob_name)
-
-        # 1. Upload flock_definition (overwrite in case of retry)
-        await blob_client.upload_blob(config.flock_definition,
-                                      overwrite=True,
-                                      content_type="text/yaml")
-        logger.debug("Uploaded blob '%s' (%d bytes)",
-                     blob_name, len(config.flock_definition.encode()))
-
-        # 2. Persist lightweight record in the table
-        tbl_client = self.table_svc.get_table_client(self._TABLE_NAME)
-        entity = {
-            "PartitionKey": self._PARTITION_KEY,
-            "RowKey":       config.share_id,
-            "agent_name":   config.agent_name,
-            "created_at":   config.created_at.isoformat(),
-            "share_type":   config.share_type,
-            "chat_message_key":  config.chat_message_key,
-            "chat_history_key":  config.chat_history_key,
-            "chat_response_key": config.chat_response_key,
-            # NEW – just a few bytes, well under 64 KiB
-            "flock_blob_name": blob_name,
-        }
-        await tbl_client.upsert_entity(entity)
-        logger.info("Saved shared link %s → blob '%s'", config.share_id, blob_name)
-        return config
-
-    # -------------------------------------------------------------- get_config
-    async def get_config(self, share_id: str) -> SharedLinkConfig | None:
-        tbl_client = self.table_svc.get_table_client(self._TABLE_NAME)
-        try:
-            entity = await tbl_client.get_entity(self._PARTITION_KEY, share_id)
-        except ResourceNotFoundError:
-            logger.debug("No config entity for id '%s'", share_id)
-            return None
-
-        blob_name   = entity["flock_blob_name"]
-        blob_client = self.blob_svc.get_blob_client(self._CONTAINER_NAME, blob_name)
-        try:
-            blob_bytes = await (await blob_client.download_blob()).readall()
-            flock_yaml = blob_bytes.decode()
-        except Exception as e:
-            logger.error("Cannot download blob '%s' for share_id=%s: %s",
-                         blob_name, share_id, e, exc_info=True)
-            raise
-
-        return SharedLinkConfig(
-            share_id          = share_id,
-            agent_name        = entity["agent_name"],
-            created_at        = entity["created_at"],
-            flock_definition  = flock_yaml,
-            share_type        = entity.get("share_type", "agent_run"),
-            chat_message_key  = entity.get("chat_message_key"),
-            chat_history_key  = entity.get("chat_history_key"),
-            chat_response_key = entity.get("chat_response_key"),
-        )
-
-    # ----------------------------------------------------------- delete_config
-    async def delete_config(self, share_id: str) -> bool:
-        tbl_client = self.table_svc.get_table_client(self._TABLE_NAME)
-        try:
-            entity = await tbl_client.get_entity(self._PARTITION_KEY, share_id)
-        except ResourceNotFoundError:
-            logger.info("Delete: entity %s not found", share_id)
-            return False
-
-        # 1. Remove blob (ignore missing blob)
-        blob_name   = entity["flock_blob_name"]
-        blob_client = self.blob_svc.get_blob_client(self._CONTAINER_NAME, blob_name)
-        try:
-            await blob_client.delete_blob(delete_snapshots="include")
-            logger.debug("Deleted blob '%s'", blob_name)
-        except ResourceNotFoundError:
-            logger.warning("Blob '%s' already gone", blob_name)
-
-        # 2. Remove table row
-        await tbl_client.delete_entity(self._PARTITION_KEY, share_id)
-        logger.info("Deleted shared link %s and its blob", share_id)
-        return True
-
-    # -------------------------------------------------------- save_feedback --
-    async def save_feedback(self, record: FeedbackRecord) -> FeedbackRecord:
-        """Persist a feedback record. If a flock_definition is present, upload it as a blob and store only a reference in the table row to avoid oversized entities (64 KiB limit)."""
-        tbl_client = self.table_svc.get_table_client(self._FEEDBACK_TBL_NAME)
-
-        # Core entity fields (avoid dumping the full Pydantic model – too many columns / large value)
-        entity: dict[str, Any] = {
-            "PartitionKey": "feedback",
-            "RowKey":       record.feedback_id,
-            "share_id":     record.share_id,
-            "context_type": record.context_type,
-            "reason":       record.reason,
-            "expected_response": record.expected_response,
-            "actual_response":   record.actual_response,
-            "created_at":   record.created_at.isoformat(),
-        }
-
-        # additional sanity check
-
-        if record.flock_name is not None:
-            entity["flock_name"] = record.flock_name
-        if record.agent_name is not None:
-            entity["agent_name"] = record.agent_name
-
-        # ------------------------------------------------------------------ YAML → Blob
-        if record.flock_definition:
-            blob_name   = f"{record.feedback_id}.yaml"
-            blob_client = self.blob_svc.get_blob_client(self._CONTAINER_NAME, blob_name)
-            # Overwrite=true so repeated feedback_id uploads (shouldn't happen) won't error
-            await blob_client.upload_blob(record.flock_definition,
-                                          overwrite=True,
-                                          content_type="text/yaml")
-            entity["flock_blob_name"] = blob_name  # lightweight reference only
-
-        # ------------------------------------------------------------------ Table upsert
-        await tbl_client.upsert_entity(entity)
-        logger.info("Saved feedback %s%s",
-                    record.feedback_id,
-                    f" → blob '{entity['flock_blob_name']}'" if "flock_blob_name" in entity else "")
-        return record
-
-    # -------------------------------------------------------- get_feedback --
-    async def get_feedback(self, id: str) -> FeedbackRecord | None:
-        """Retrieve a single feedback record from Azure Table Storage."""
-        tbl_client = self.table_svc.get_table_client(self._FEEDBACK_TBL_NAME)
-        try:
-            entity = await tbl_client.get_entity("feedback", id)
-        except ResourceNotFoundError:
-            logger.debug("No feedback record found for ID: %s", id)
-            return None
-
-        # Get flock_definition from blob if it exists
-        flock_definition = None
-        if "flock_blob_name" in entity:
-            blob_name = entity["flock_blob_name"]
-            blob_client = self.blob_svc.get_blob_client(self._CONTAINER_NAME, blob_name)
-            try:
-                blob_bytes = await (await blob_client.download_blob()).readall()
-                flock_definition = blob_bytes.decode()
-            except Exception as e:
-                logger.error("Cannot download blob '%s' for feedback_id=%s: %s",
-                           blob_name, id, e, exc_info=True)
-                # Continue without flock_definition rather than failing
-
-        return FeedbackRecord(
-            feedback_id=id,
-            share_id=entity.get("share_id"),
-            context_type=entity["context_type"],
-            reason=entity["reason"],
-            expected_response=entity.get("expected_response"),
-            actual_response=entity.get("actual_response"),
-            flock_name=entity.get("flock_name"),
-            agent_name=entity.get("agent_name"),
-            flock_definition=flock_definition,
-            created_at=entity["created_at"],
-        )
-
-
-    # ------------------------------------------------ get_all_feedback_records --
-    async def get_all_feedback_records_for_agent(self, agent_name: str) -> list[FeedbackRecord]:
-        """Retrieve all feedback records from Azure Table Storage for a specific agent."""
-        tbl_client = self.table_svc.get_table_client(self._FEEDBACK_TBL_NAME)
-
-        # Use Azure Table Storage filtering to only get records for the specified agent
-        escaped_agent_name = agent_name.replace("'", "''")
-        filter_query = f"agent_name eq '{escaped_agent_name}'"
-
-        logger.debug(f"Querying feedback records with filter: {filter_query}")
-
-        records = []
-        try:
-            async for entity in tbl_client.query_entities(filter_query):
-                # Get flock_definition from blob if it exists
-                flock_definition = None
-                if "flock_blob_name" in entity:
-                    blob_name = entity["flock_blob_name"]
-                    blob_client = self.blob_svc.get_blob_client(self._CONTAINER_NAME, blob_name)
-                    try:
-                        blob_bytes = await (await blob_client.download_blob()).readall()
-                        flock_definition = blob_bytes.decode()
-                    except Exception as e:
-                        logger.error("Cannot download blob '%s' for feedback_id=%s: %s",
-                                   blob_name, entity["RowKey"], e, exc_info=True)
-                        # Continue without flock_definition rather than failing
-
-                records.append(FeedbackRecord(
-                    feedback_id=entity["RowKey"],
-                    share_id=entity.get("share_id"),
-                    context_type=entity["context_type"],
-                    reason=entity["reason"],
-                    expected_response=entity.get("expected_response"),
-                    actual_response=entity.get("actual_response"),
-                    flock_name=entity.get("flock_name"),
-                    agent_name=entity.get("agent_name"),
-                    flock_definition=flock_definition,
-                    created_at=entity["created_at"],
-                ))
-
-            logger.debug("Retrieved %d feedback records for agent %s", len(records), agent_name)
-            return records
-
-        except Exception as e:
-            # Log the error.
-            logger.error(
-                f"Unable to query entries for agent {agent_name}. Exception: {e}"
-            )
-            return records
-
-
-# ----------------------- Factory Function -----------------------
-
-def create_shared_link_store(store_type: str | None = None, connection_string: str | None = None) -> SharedLinkStoreInterface:
-    """Factory function to create the appropriate shared link store based on configuration.
-
-    Args:
-        store_type: Type of store to create ("local" for SQLite, "azure-storage" for Azure Table Storage)
-        connection_string: Connection string for the store (file path for SQLite, connection string for Azure)
-
-    Returns:
-        Configured SharedLinkStoreInterface implementation
-    """
-    import os
-
-    # Get values from environment if not provided
-    if store_type is None:
-        store_type = os.getenv("FLOCK_WEBAPP_STORE", "local").lower()
-
-    if connection_string is None:
-        connection_string = os.getenv("FLOCK_WEBAPP_STORE_CONNECTION", ".flock/shared_links.db")
-
-    if store_type == "local":
-        return SQLiteSharedLinkStore(connection_string)
-    elif store_type == "azure-storage":
-        return AzureTableSharedLinkStore(connection_string)
-    else:
-        raise ValueError(f"Unsupported store type: {store_type}. Supported types: 'local', 'azure-storage'")