agent-platform/services/knowledge-service/app/schemas/knowledge.py

from typing import TYPE_CHECKING

from core_domain import (
    KnowledgeBaseContract,
    KnowledgeBaseStatus,
    KnowledgeChunkContract,
    KnowledgeDocumentContract,
    KnowledgeSearchRequestContract,
    KnowledgeSearchResultContract,
)
from core_shared import JSONValue
from pydantic import BaseModel, Field

if TYPE_CHECKING:
    from app.db.models import KnowledgeBase, KnowledgeChunk, KnowledgeDocument


class KnowledgeBaseCreateRequest(BaseModel):
    code: str
    name: str
    description: str | None = None
    metadata_json: dict[str, JSONValue] = Field(default_factory=dict)


class KnowledgeBaseStatusUpdateRequest(BaseModel):
    status: KnowledgeBaseStatus


class KnowledgeBaseResponse(KnowledgeBaseContract):
    @classmethod
    def from_entity(cls, entity: "KnowledgeBase") -> "KnowledgeBaseResponse":
        return cls.model_validate(entity, from_attributes=True)


class KnowledgeDocumentCreateRequest(BaseModel):
    knowledge_base_id: str
    title: str
    content_text: str | None = None
    content_base64: str | None = None
    source_type: str = "text"
    source_uri: str | None = None
    metadata_json: dict[str, JSONValue] = Field(default_factory=dict)
    chunk_size: int | None = Field(default=None, gt=0)
    chunk_overlap: int | None = Field(default=None, ge=0)


class KnowledgeDocumentResponse(KnowledgeDocumentContract):
    @classmethod
    def from_entity(cls, entity: "KnowledgeDocument") -> "KnowledgeDocumentResponse":
        return cls.model_validate(entity, from_attributes=True)


class KnowledgeChunkResponse(KnowledgeChunkContract):
    @classmethod
    def from_entity(cls, entity: "KnowledgeChunk") -> "KnowledgeChunkResponse":
        return cls.model_validate(entity, from_attributes=True)


class KnowledgeDocumentIngestResponse(BaseModel):
    document: KnowledgeDocumentResponse
    chunks: list[KnowledgeChunkResponse]


class KnowledgeDocumentParseRequest(BaseModel):
    source_type: str = "auto"
    source_uri: str | None = None
    content_text: str | None = None
    content_base64: str | None = None


class KnowledgeDocumentParseResponse(BaseModel):
    content_text: str
    source_type: str
    metadata_json: dict[str, JSONValue] = Field(default_factory=dict)


class KnowledgeSearchRequest(KnowledgeSearchRequestContract):
    pass


class KnowledgeSearchResultResponse(KnowledgeSearchResultContract):
    pass