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

from typing import TYPE_CHECKING

from pydantic import BaseModel

from core_domain import (
    MemoryCreateContract,
    MemoryItemContract,
    MemorySearchRequestContract,
    MemorySearchResultContract,
    MemoryStatus,
)
from core_shared import JSONValue

if TYPE_CHECKING:
    from app.db.models import MemoryItem


class MemoryCreateRequest(MemoryCreateContract):
    pass


class MemoryStatusUpdateRequest(BaseModel):
    tenant_id: str
    status: MemoryStatus


class MemorySearchRequest(MemorySearchRequestContract):
    pass


class MemoryResponse(MemoryItemContract):
    @classmethod
    def from_entity(cls, entity: "MemoryItem") -> "MemoryResponse":
        return cls.model_validate(entity, from_attributes=True)


class MemorySearchResultResponse(MemorySearchResultContract):
    item: MemoryResponse

    @classmethod
    def from_entity(
        cls,
        entity: "MemoryItem",
        score: float,
        score_json: dict[str, JSONValue],
    ) -> "MemorySearchResultResponse":
        return cls(
            item=MemoryResponse.from_entity(entity),
            score=score,
            score_json=score_json,
        )