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

from typing import TYPE_CHECKING

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

if TYPE_CHECKING:
    from app.db.models import MemoryItem


class MemoryCreateRequest(MemoryCreateContract):
    pass


class MemoryStatusUpdateRequest(BaseModel):
    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)