llm-in-text/backend/docs_store.py

import mimetypes
import os
import threading
import uuid
from dataclasses import dataclass
from typing import Any

try:  # pragma: no cover - optional in some test paths
    import psycopg
    from psycopg.rows import dict_row
except Exception:  # pragma: no cover
    psycopg = None
    dict_row = None


MAX_TEXT_SIZE = 8 * 1024 * 1024
PREVIEW_TEXT_SIZE = 2 * 1024 * 1024

TEXT_EXTENSIONS = {
    "md", "markdown", "txt", "json", "js", "jsx", "ts", "tsx",
    "css", "scss", "less", "html", "htm", "py", "vue", "xml",
    "yaml", "yml", "csv", "log", "sql", "toml", "ini", "cfg",
    "conf", "sh", "bat", "ps1", "java", "c", "cpp", "h", "hpp",
    "go", "rs", "swift", "kt", "rb", "php", "pl", "r", "scala",
    "gradle", "properties", "env", "gitignore", "dockerfile",
}

BINARY_EXTENSIONS = {
    "exe", "dll", "so", "dylib", "bin", "dat", "obj", "o", "a",
    "doc", "docx", "xls", "xlsx", "ppt", "pptx", "odt", "ods", "odp",
    "pdf", "zip", "rar", "7z", "tar", "gz", "bz2", "xz",
    "png", "jpg", "jpeg", "gif", "bmp", "ico", "webp", "svg",
    "mp3", "mp4", "wav", "avi", "mov", "mkv", "flv", "wmv",
    "ttf", "otf", "woff", "woff2", "eot",
    "class", "pyc", "pyo", "jar", "war", "ear",
    "db", "sqlite", "mdb", "accdb",
    "pem", "key", "crt", "cer", "p12", "pfx", "jks",
    "msg", "eml", "pst", "ost",
    "dwg", "dxf", "step", "stl", "fbx", "3ds", "blend",
}

DEFAULT_MIME_TYPES = {
    "avi": "video/x-msvideo",
    "md": "text/markdown",
    "markdown": "text/markdown",
    "mkv": "video/x-matroska",
    "mov": "video/quicktime",
    "mp4": "video/mp4",
    "txt": "text/plain",
    "json": "application/json",
    "js": "text/javascript",
    "jsx": "text/javascript",
    "ts": "text/typescript",
    "tsx": "text/typescript",
    "css": "text/css",
    "html": "text/html",
    "htm": "text/html",
    "py": "text/x-python",
    "vue": "text/plain",
    "xml": "application/xml",
    "yaml": "text/yaml",
    "yml": "text/yaml",
    "csv": "text/csv",
    "log": "text/plain",
    "sql": "text/plain",
    "toml": "text/plain",
    "ini": "text/plain",
    "cfg": "text/plain",
    "conf": "text/plain",
    "sh": "text/plain",
    "bat": "text/plain",
    "ps1": "text/plain",
    "jpg": "image/jpeg",
    "jpeg": "image/jpeg",
    "flv": "video/x-flv",
    "m4v": "video/x-m4v",
    "png": "image/png",
    "gif": "image/gif",
    "webp": "image/webp",
    "svg": "image/svg+xml",
    "pdf": "application/pdf",
    "ogg": "video/ogg",
    "ogv": "video/ogg",
    "webm": "video/webm",
    "wmv": "video/x-ms-wmv",
}


def _now_sql() -> str:
    return "CURRENT_TIMESTAMP"


def get_extension(name: str = "") -> str:
    value = str(name or "")
    if "." not in value:
        return value.lower() if value.lower() == "dockerfile" else ""
    return value.rsplit(".", 1)[-1].lower()


def infer_mime_type(name: str, fallback: str = "") -> str:
    if fallback:
        return fallback
    ext = get_extension(name)
    guessed = DEFAULT_MIME_TYPES.get(ext)
    if guessed:
        return guessed
    guessed, _ = mimetypes.guess_type(name)
    return guessed or "application/octet-stream"


def is_text_file(name: str, mime_type: str = "") -> bool:
    ext = get_extension(name)
    if ext in BINARY_EXTENSIONS:
        return False
    mime_value = str(mime_type or "").lower()
    return ext in TEXT_EXTENSIONS or mime_value.startswith("text/") or "json" in mime_value or "xml" in mime_value


def _text_payload(text: str) -> dict[str, Any]:
    encoded = text.encode("utf-8")
    preview = encoded[:PREVIEW_TEXT_SIZE].decode("utf-8", errors="ignore")
    return {
        "storage_kind": "text",
        "size": len(encoded),
        "content_text": text,
        "preview_text": preview if len(encoded) > PREVIEW_TEXT_SIZE else text,
        "is_truncated_preview": len(encoded) > PREVIEW_TEXT_SIZE,
        "blob_data": None,
    }


def prepare_file_payload(name: str, raw_bytes: bytes, mime_type: str = "") -> dict[str, Any]:
    resolved_mime = infer_mime_type(name, mime_type)
    if is_text_file(name, resolved_mime):
        text = raw_bytes.decode("utf-8", errors="ignore")
        return {
            "mime_type": resolved_mime,
            **_text_payload(text),
        }
    return {
        "mime_type": resolved_mime,
        "storage_kind": "blob",
        "size": len(raw_bytes),
        "content_text": "",
        "preview_text": "",
        "is_truncated_preview": False,
        "blob_data": raw_bytes,
    }


@dataclass
class BlobPayload:
    content: bytes
    mime_type: str
    filename: str


class BaseDocumentStore:
    _UNSET = object()

    def list_nodes(self) -> list[dict[str, Any]]:
        raise NotImplementedError

    def create_folder(self, name: str, parent_id: str | None) -> dict[str, Any]:
        raise NotImplementedError

    def create_text_file(self, name: str, parent_id: str | None, content: str = "") -> dict[str, Any]:
        raise NotImplementedError

    def upload_file(self, name: str, parent_id: str | None, raw_bytes: bytes, mime_type: str = "") -> dict[str, Any]:
        raise NotImplementedError

    def update_node(self, node_id: str, *, name: str | None | object = _UNSET, parent_id: str | None | object = _UNSET, content: str | None | object = _UNSET) -> dict[str, Any]:
        raise NotImplementedError

    def replace_blob(self, node_id: str, filename: str, raw_bytes: bytes, mime_type: str = "") -> dict[str, Any]:
        raise NotImplementedError

    def delete_node(self, node_id: str) -> None:
        raise NotImplementedError

    def get_blob(self, node_id: str) -> BlobPayload:
        raise NotImplementedError


class InMemoryDocumentStore(BaseDocumentStore):
    def __init__(self) -> None:
        self.nodes: dict[str, dict[str, Any]] = {}

    def _serialize(self, node: dict[str, Any]) -> dict[str, Any]:
        return {
            "id": node["id"],
            "parentId": node["parent_id"],
            "type": node["type"],
            "name": node["name"],
            "mimeType": node.get("mime_type") or "",
            "storageKind": node.get("storage_kind") or "text",
            "size": int(node.get("size") or 0),
            "content": node.get("content_text") or "",
            "previewText": node.get("preview_text") or "",
            "isTruncatedPreview": bool(node.get("is_truncated_preview")),
            "createdAt": int(node["created_at"]),
            "updatedAt": int(node["updated_at"]),
        }

    def _timestamp(self) -> int:
        import time
        return int(time.time() * 1000)

    def _get(self, node_id: str) -> dict[str, Any]:
        node = self.nodes.get(node_id)
        if node is None:
            raise KeyError(node_id)
        return node

    def list_nodes(self) -> list[dict[str, Any]]:
        ordered = sorted(
            self.nodes.values(),
            key=lambda item: (item["type"] != "folder", item["name"].lower(), item["created_at"]),
        )
        return [self._serialize(node) for node in ordered]

    def create_folder(self, name: str, parent_id: str | None) -> dict[str, Any]:
        now = self._timestamp()
        node = {
            "id": str(uuid.uuid4()),
            "parent_id": parent_id,
            "type": "folder",
            "name": name,
            "mime_type": "",
            "storage_kind": "text",
            "size": 0,
            "content_text": "",
            "preview_text": "",
            "is_truncated_preview": False,
            "blob_data": None,
            "created_at": now,
            "updated_at": now,
        }
        self.nodes[node["id"]] = node
        return self._serialize(node)

    def create_text_file(self, name: str, parent_id: str | None, content: str = "") -> dict[str, Any]:
        now = self._timestamp()
        payload = _text_payload(content)
        node = {
            "id": str(uuid.uuid4()),
            "parent_id": parent_id,
            "type": "file",
            "name": name,
            "mime_type": infer_mime_type(name),
            **payload,
            "created_at": now,
            "updated_at": now,
        }
        self.nodes[node["id"]] = node
        return self._serialize(node)

    def upload_file(self, name: str, parent_id: str | None, raw_bytes: bytes, mime_type: str = "") -> dict[str, Any]:
        now = self._timestamp()
        payload = prepare_file_payload(name, raw_bytes, mime_type)
        node = {
            "id": str(uuid.uuid4()),
            "parent_id": parent_id,
            "type": "file",
            "name": name,
            **payload,
            "created_at": now,
            "updated_at": now,
        }
        self.nodes[node["id"]] = node
        return self._serialize(node)

    def update_node(self, node_id: str, *, name: str | None | object = BaseDocumentStore._UNSET, parent_id: str | None | object = BaseDocumentStore._UNSET, content: str | None | object = BaseDocumentStore._UNSET) -> dict[str, Any]:
        node = self._get(node_id)
        if name is not BaseDocumentStore._UNSET:
            node["name"] = name
            if node["type"] == "file":
                node["mime_type"] = infer_mime_type(name, node.get("mime_type") or "")
        if parent_id is not BaseDocumentStore._UNSET:
            node["parent_id"] = parent_id
        if content is not BaseDocumentStore._UNSET:
            payload = _text_payload(content)
            node.update(payload)
        node["updated_at"] = self._timestamp()
        return self._serialize(node)

    def replace_blob(self, node_id: str, filename: str, raw_bytes: bytes, mime_type: str = "") -> dict[str, Any]:
        node = self._get(node_id)
        payload = prepare_file_payload(filename, raw_bytes, mime_type)
        node.update(payload)
        node["name"] = filename
        node["mime_type"] = payload["mime_type"]
        node["updated_at"] = self._timestamp()
        return self._serialize(node)

    def delete_node(self, node_id: str) -> None:
        descendants = {node_id}
        changed = True
        while changed:
            changed = False
            for node in list(self.nodes.values()):
                if node.get("parent_id") in descendants and node["id"] not in descendants:
                    descendants.add(node["id"])
                    changed = True
        for current_id in descendants:
            self.nodes.pop(current_id, None)

    def get_blob(self, node_id: str) -> BlobPayload:
        node = self._get(node_id)
        if node["type"] != "file":
            raise FileNotFoundError(node_id)
        if node.get("blob_data") is not None:
            content = node["blob_data"]
        else:
            content = (node.get("content_text") or "").encode("utf-8")
        return BlobPayload(
            content=content,
            mime_type=node.get("mime_type") or infer_mime_type(node.get("name") or ""),
            filename=node["name"],
        )


class PostgresDocumentStore(BaseDocumentStore):
    def __init__(self, database_url: str) -> None:
        if psycopg is None or dict_row is None:
            raise RuntimeError("psycopg 未安装，无法使用 PostgreSQL 文档存储")
        self.database_url = database_url
        self._init_lock = threading.Lock()
        self._initialized = False

    def _connect(self):
        return psycopg.connect(self.database_url, autocommit=True, row_factory=dict_row)

    def _ensure_initialized(self) -> None:
        if self._initialized:
            return
        with self._init_lock:
            if self._initialized:
                return
            with self._connect() as conn:
                with conn.cursor() as cur:
                    cur.execute(
                        """
                        CREATE TABLE IF NOT EXISTS document_nodes (
                            id TEXT PRIMARY KEY,
                            parent_id TEXT REFERENCES document_nodes(id) ON DELETE CASCADE,
                            type TEXT NOT NULL CHECK (type IN ('folder', 'file')),
                            name TEXT NOT NULL,
                            mime_type TEXT NOT NULL DEFAULT '',
                            storage_kind TEXT NOT NULL DEFAULT 'text',
                            size BIGINT NOT NULL DEFAULT 0,
                            content_text TEXT NOT NULL DEFAULT '',
                            preview_text TEXT NOT NULL DEFAULT '',
                            is_truncated_preview BOOLEAN NOT NULL DEFAULT FALSE,
                            blob_data BYTEA,
                            created_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP,
                            updated_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP
                        )
                        """
                    )
                    cur.execute(
                        "CREATE INDEX IF NOT EXISTS document_nodes_parent_idx ON document_nodes(parent_id)"
                    )
            self._initialized = True

    def _serialize(self, row: dict[str, Any]) -> dict[str, Any]:
        created = row["created_at"]
        updated = row["updated_at"]
        return {
            "id": row["id"],
            "parentId": row["parent_id"],
            "type": row["type"],
            "name": row["name"],
            "mimeType": row.get("mime_type") or "",
            "storageKind": row.get("storage_kind") or "text",
            "size": int(row.get("size") or 0),
            "content": row.get("content_text") or "",
            "previewText": row.get("preview_text") or "",
            "isTruncatedPreview": bool(row.get("is_truncated_preview")),
            "createdAt": int(created.timestamp() * 1000),
            "updatedAt": int(updated.timestamp() * 1000),
        }

    def _fetch_node(self, node_id: str) -> dict[str, Any]:
        self._ensure_initialized()
        with self._connect() as conn:
            with conn.cursor() as cur:
                cur.execute("SELECT * FROM document_nodes WHERE id = %s", (node_id,))
                row = cur.fetchone()
                if row is None:
                    raise KeyError(node_id)
                return row

    def list_nodes(self) -> list[dict[str, Any]]:
        self._ensure_initialized()
        with self._connect() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    """
                    SELECT
                        id,
                        parent_id,
                        type,
                        name,
                        mime_type,
                        storage_kind,
                        size,
                        CASE
                            WHEN storage_kind = 'text' AND size <= %s THEN content_text
                            ELSE ''
                        END AS content_text,
                        preview_text,
                        is_truncated_preview,
                        created_at,
                        updated_at
                    FROM document_nodes
                    ORDER BY
                        CASE WHEN type = 'folder' THEN 0 ELSE 1 END,
                        LOWER(name),
                        created_at
                    """,
                    (MAX_TEXT_SIZE,),
                )
                return [self._serialize(row) for row in cur.fetchall()]

    def create_folder(self, name: str, parent_id: str | None) -> dict[str, Any]:
        self._ensure_initialized()
        node_id = str(uuid.uuid4())
        with self._connect() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    """
                    INSERT INTO document_nodes (id, parent_id, type, name, updated_at)
                    VALUES (%s, %s, 'folder', %s, CURRENT_TIMESTAMP)
                    RETURNING *
                    """,
                    (node_id, parent_id, name),
                )
                return self._serialize(cur.fetchone())

    def create_text_file(self, name: str, parent_id: str | None, content: str = "") -> dict[str, Any]:
        self._ensure_initialized()
        node_id = str(uuid.uuid4())
        payload = _text_payload(content)
        with self._connect() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    """
                    INSERT INTO document_nodes (
                        id, parent_id, type, name, mime_type, storage_kind, size,
                        content_text, preview_text, is_truncated_preview, blob_data, updated_at
                    )
                    VALUES (%s, %s, 'file', %s, %s, %s, %s, %s, %s, %s, %s, CURRENT_TIMESTAMP)
                    RETURNING *
                    """,
                    (
                        node_id,
                        parent_id,
                        name,
                        infer_mime_type(name),
                        payload["storage_kind"],
                        payload["size"],
                        payload["content_text"],
                        payload["preview_text"],
                        payload["is_truncated_preview"],
                        None,
                    ),
                )
                return self._serialize(cur.fetchone())

    def upload_file(self, name: str, parent_id: str | None, raw_bytes: bytes, mime_type: str = "") -> dict[str, Any]:
        self._ensure_initialized()
        node_id = str(uuid.uuid4())
        payload = prepare_file_payload(name, raw_bytes, mime_type)
        with self._connect() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    """
                    INSERT INTO document_nodes (
                        id, parent_id, type, name, mime_type, storage_kind, size,
                        content_text, preview_text, is_truncated_preview, blob_data, updated_at
                    )
                    VALUES (%s, %s, 'file', %s, %s, %s, %s, %s, %s, %s, %s, CURRENT_TIMESTAMP)
                    RETURNING *
                    """,
                    (
                        node_id,
                        parent_id,
                        name,
                        payload["mime_type"],
                        payload["storage_kind"],
                        payload["size"],
                        payload["content_text"],
                        payload["preview_text"],
                        payload["is_truncated_preview"],
                        payload["blob_data"],
                    ),
                )
                return self._serialize(cur.fetchone())

    def update_node(self, node_id: str, *, name: str | None | object = BaseDocumentStore._UNSET, parent_id: str | None | object = BaseDocumentStore._UNSET, content: str | None | object = BaseDocumentStore._UNSET) -> dict[str, Any]:
        self._ensure_initialized()
        row = self._fetch_node(node_id)
        next_name = row["name"] if name is BaseDocumentStore._UNSET else name
        next_parent_id = row["parent_id"] if parent_id is BaseDocumentStore._UNSET else parent_id
        next_mime_type = row.get("mime_type") or ""
        next_storage_kind = row.get("storage_kind") or "text"
        next_size = row.get("size") or 0
        next_content_text = row.get("content_text") or ""
        next_preview_text = row.get("preview_text") or ""
        next_is_truncated = bool(row.get("is_truncated_preview"))
        next_blob_data = row.get("blob_data")

        if row["type"] == "file" and name is not BaseDocumentStore._UNSET:
            next_mime_type = infer_mime_type(next_name, next_mime_type)
        if content is not BaseDocumentStore._UNSET:
            payload = _text_payload(content)
            next_storage_kind = payload["storage_kind"]
            next_size = payload["size"]
            next_content_text = payload["content_text"]
            next_preview_text = payload["preview_text"]
            next_is_truncated = payload["is_truncated_preview"]
            next_blob_data = None

        with self._connect() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    """
                    UPDATE document_nodes
                    SET
                        name = %s,
                        parent_id = %s,
                        mime_type = %s,
                        storage_kind = %s,
                        size = %s,
                        content_text = %s,
                        preview_text = %s,
                        is_truncated_preview = %s,
                        blob_data = %s,
                        updated_at = CURRENT_TIMESTAMP
                    WHERE id = %s
                    RETURNING *
                    """,
                    (
                        next_name,
                        next_parent_id,
                        next_mime_type,
                        next_storage_kind,
                        next_size,
                        next_content_text,
                        next_preview_text,
                        next_is_truncated,
                        next_blob_data,
                        node_id,
                    ),
                )
                updated = cur.fetchone()
                if updated is None:
                    raise KeyError(node_id)
                return self._serialize(updated)

    def replace_blob(self, node_id: str, filename: str, raw_bytes: bytes, mime_type: str = "") -> dict[str, Any]:
        self._ensure_initialized()
        payload = prepare_file_payload(filename, raw_bytes, mime_type)
        with self._connect() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    """
                    UPDATE document_nodes
                    SET
                        name = %s,
                        mime_type = %s,
                        storage_kind = %s,
                        size = %s,
                        content_text = %s,
                        preview_text = %s,
                        is_truncated_preview = %s,
                        blob_data = %s,
                        updated_at = CURRENT_TIMESTAMP
                    WHERE id = %s AND type = 'file'
                    RETURNING *
                    """,
                    (
                        filename,
                        payload["mime_type"],
                        payload["storage_kind"],
                        payload["size"],
                        payload["content_text"],
                        payload["preview_text"],
                        payload["is_truncated_preview"],
                        payload["blob_data"],
                        node_id,
                    ),
                )
                updated = cur.fetchone()
                if updated is None:
                    raise KeyError(node_id)
                return self._serialize(updated)

    def delete_node(self, node_id: str) -> None:
        self._ensure_initialized()
        with self._connect() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    """
                    WITH RECURSIVE descendants AS (
                        SELECT id FROM document_nodes WHERE id = %s
                        UNION ALL
                        SELECT child.id
                        FROM document_nodes child
                        INNER JOIN descendants parent ON child.parent_id = parent.id
                    )
                    DELETE FROM document_nodes
                    WHERE id IN (SELECT id FROM descendants)
                    """,
                    (node_id,),
                )

    def get_blob(self, node_id: str) -> BlobPayload:
        self._ensure_initialized()
        with self._connect() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    """
                    SELECT id, name, mime_type, content_text, blob_data
                    FROM document_nodes
                    WHERE id = %s AND type = 'file'
                    """,
                    (node_id,),
                )
                row = cur.fetchone()
                if row is None:
                    raise FileNotFoundError(node_id)
                blob_data = row.get("blob_data")
                if blob_data is None:
                    content = (row.get("content_text") or "").encode("utf-8")
                else:
                    content = bytes(blob_data)
                return BlobPayload(
                    content=content,
                    mime_type=row.get("mime_type") or infer_mime_type(row.get("name") or ""),
                    filename=row["name"],
                )


_document_store: BaseDocumentStore | None = None


def get_document_store() -> BaseDocumentStore:
    global _document_store
    if _document_store is not None:
        return _document_store
    backend = (os.getenv("DOCS_BACKEND") or "postgres").strip().lower()
    if backend == "memory":
        _document_store = InMemoryDocumentStore()
        return _document_store

    database_url = os.getenv("DATABASE_URL", "").strip()
    if not database_url:
        raise RuntimeError("缺少 DATABASE_URL，无法初始化 PostgreSQL 文档存储")
    _document_store = PostgresDocumentStore(database_url)
    return _document_store


def reset_document_store() -> None:
    global _document_store
    _document_store = None