agent-platform/services/knowledge-service/app/domain/repositories.py

from datetime import datetime

from sqlalchemy import delete, select
from sqlalchemy.orm import Session

from core_domain import KnowledgeBaseStatus, KnowledgeDocumentStatus
from core_shared import JSONValue

from app.db.models import KnowledgeBase, KnowledgeChunk, KnowledgeDocument


class KnowledgeBaseRepository:
    def __init__(self, db: Session) -> None:
        self.db = db

    def create(
        self,
        *,
        tenant_id: str,
        code: str,
        name: str,
        description: str | None,
        metadata_json: dict[str, JSONValue] | None,
    ) -> KnowledgeBase:
        entity = KnowledgeBase(
            tenant_id=tenant_id,
            code=code,
            name=name,
            description=description,
            metadata_json=metadata_json,
        )
        self.db.add(entity)
        self.db.commit()
        self.db.refresh(entity)
        return entity

    def list_by_tenant(self, *, tenant_id: str) -> list[KnowledgeBase]:
        stmt = (
            select(KnowledgeBase)
            .where(KnowledgeBase.tenant_id == tenant_id)
            .order_by(KnowledgeBase.created_time.desc())
        )
        return list(self.db.scalars(stmt))

    def get_by_id(self, *, tenant_id: str, knowledge_base_id: str) -> KnowledgeBase | None:
        stmt = (
            select(KnowledgeBase)
            .where(KnowledgeBase.tenant_id == tenant_id)
            .where(KnowledgeBase.id == knowledge_base_id)
        )
        return self.db.scalar(stmt)

    def update_status(
        self,
        *,
        tenant_id: str,
        knowledge_base_id: str,
        status: KnowledgeBaseStatus,
    ) -> KnowledgeBase | None:
        entity = self.get_by_id(tenant_id=tenant_id, knowledge_base_id=knowledge_base_id)
        if entity is None:
            return None
        entity.status = status
        self.db.commit()
        self.db.refresh(entity)
        return entity


class KnowledgeDocumentRepository:
    def __init__(self, db: Session) -> None:
        self.db = db

    def create(
        self,
        *,
        tenant_id: str,
        knowledge_base_id: str,
        title: str,
        source_type: str,
        source_uri: str | None,
        content_text: str,
        content_hash: str | None,
        metadata_json: dict[str, JSONValue] | None,
    ) -> KnowledgeDocument:
        entity = KnowledgeDocument(
            tenant_id=tenant_id,
            knowledge_base_id=knowledge_base_id,
            title=title,
            source_type=source_type,
            source_uri=source_uri,
            content_text=content_text,
            content_hash=content_hash,
            metadata_json=metadata_json,
            status="draft",
        )
        self.db.add(entity)
        self.db.commit()
        self.db.refresh(entity)
        return entity

    def list_by_base(
        self,
        *,
        tenant_id: str,
        knowledge_base_id: str,
    ) -> list[KnowledgeDocument]:
        stmt = (
            select(KnowledgeDocument)
            .where(KnowledgeDocument.tenant_id == tenant_id)
            .where(KnowledgeDocument.knowledge_base_id == knowledge_base_id)
            .order_by(KnowledgeDocument.created_time.desc())
        )
        return list(self.db.scalars(stmt))

    def get_by_id(self, *, tenant_id: str, document_id: str) -> KnowledgeDocument | None:
        stmt = (
            select(KnowledgeDocument)
            .where(KnowledgeDocument.tenant_id == tenant_id)
            .where(KnowledgeDocument.id == document_id)
        )
        return self.db.scalar(stmt)

    def update_status(
        self,
        *,
        document_id: str,
        status: KnowledgeDocumentStatus,
    ) -> KnowledgeDocument | None:
        entity = self.db.get(KnowledgeDocument, document_id)
        if entity is None:
            return None
        entity.status = status
        entity.indexed_time = datetime.utcnow() if status == "indexed" else entity.indexed_time
        self.db.commit()
        self.db.refresh(entity)
        return entity


class KnowledgeChunkRepository:
    def __init__(self, db: Session) -> None:
        self.db = db

    def replace_document_chunks(
        self,
        *,
        tenant_id: str,
        knowledge_base_id: str,
        document_id: str,
        chunks: list[dict[str, JSONValue]],
    ) -> list[KnowledgeChunk]:
        self.db.execute(
            delete(KnowledgeChunk)
            .where(KnowledgeChunk.tenant_id == tenant_id)
            .where(KnowledgeChunk.document_id == document_id)
        )
        entities: list[KnowledgeChunk] = []
        for chunk in chunks:
            entity = KnowledgeChunk(
                tenant_id=tenant_id,
                knowledge_base_id=knowledge_base_id,
                document_id=document_id,
                chunk_index=_read_int(chunk, "chunk_index"),
                content_text=_read_string(chunk, "content_text"),
                token_count=_read_int(chunk, "token_count"),
                embedding_model=_read_optional_string(chunk, "embedding_model"),
                embedding_json=_read_float_list(chunk, "embedding_json"),
                metadata_json=_read_optional_dict(chunk, "metadata_json"),
            )
            self.db.add(entity)
            entities.append(entity)
        self.db.commit()
        for entity in entities:
            self.db.refresh(entity)
        return entities

    def list_by_base(
        self,
        *,
        tenant_id: str,
        knowledge_base_id: str,
    ) -> list[KnowledgeChunk]:
        stmt = (
            select(KnowledgeChunk)
            .where(KnowledgeChunk.tenant_id == tenant_id)
            .where(KnowledgeChunk.knowledge_base_id == knowledge_base_id)
            .order_by(KnowledgeChunk.created_time.asc())
        )
        return list(self.db.scalars(stmt))


def _read_string(payload: dict[str, JSONValue], key: str) -> str:
    value = payload.get(key)
    return value if isinstance(value, str) else ""


def _read_optional_string(payload: dict[str, JSONValue], key: str) -> str | None:
    value = payload.get(key)
    return value if isinstance(value, str) else None


def _read_int(payload: dict[str, JSONValue], key: str) -> int:
    value = payload.get(key)
    return value if isinstance(value, int) and not isinstance(value, bool) else 0


def _read_float_list(payload: dict[str, JSONValue], key: str) -> list[float] | None:
    value = payload.get(key)
    if not isinstance(value, list):
        return None
    return [float(item) for item in value if isinstance(item, (int, float))]


def _read_optional_dict(
    payload: dict[str, JSONValue],
    key: str,
) -> dict[str, JSONValue] | None:
    value = payload.get(key)
    if isinstance(value, dict):
        return {str(item_key): item_value for item_key, item_value in value.items()}
    return None