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

from datetime import datetime
from typing import TYPE_CHECKING, Generic, TypeVar

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

if TYPE_CHECKING:
    from app.db.models import MemoryItem

T = TypeVar("T")


class ApiErrorResponse(BaseModel):
    errorType: str
    message: str
    details: dict[str, JSONValue] = Field(default_factory=dict)


class ApiResponse(BaseModel, Generic[T]):
    success: bool = True
    data: T | None = None
    error: ApiErrorResponse | None = None
    requestId: str
    serverTime: datetime


class PageRequest(BaseModel):
    page: int = Field(default=1, ge=1)
    pageSize: int = Field(default=20, ge=1, le=200)
    keyword: str | None = None

    @property
    def offset(self) -> int:
        return (self.page - 1) * self.pageSize


class PageResult(BaseModel, Generic[T]):
    items: list[T]
    total: int
    page: int
    pageSize: int
    hasMore: bool

    @classmethod
    def from_items(
        cls,
        *,
        items: list[T],
        total: int,
        page: int,
        page_size: int) -> "PageResult[T]":
        return cls(
            items=items,
            total=total,
            page=page,
            pageSize=page_size,
            hasMore=page * page_size < total)


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)


class MemoryItemDto(BaseModel):
    id: str
    scopeType: MemoryScopeType
    scopeId: str
    memoryType: str
    contentText: str
    content: dict[str, JSONValue] | None = None
    metadata: dict[str, JSONValue] = Field(default_factory=dict)
    embeddingModel: str | None = None
    embedding: list[float] | None = None
    ownerAgentId: str | None = None
    userId: str | None = None
    sessionId: str | None = None
    sourceRef: str | None = None
    importanceScore: int
    status: MemoryStatus
    lastAccessedTime: datetime | None = None
    expiresTime: datetime | None = None
    createdTime: datetime

    @classmethod
    def from_entity(cls, entity: "MemoryItem") -> "MemoryItemDto":
        return cls(
            id=entity.id,
            scopeType=entity.scope_type,
            scopeId=entity.scope_id,
            memoryType=entity.memory_type,
            contentText=entity.content_text,
            content=entity.content_json,
            metadata=entity.metadata_json or {},
            embeddingModel=entity.embedding_model,
            embedding=entity.embedding_json,
            ownerAgentId=entity.owner_agent_id,
            userId=entity.user_id,
            sessionId=entity.session_id,
            sourceRef=entity.source_ref,
            importanceScore=entity.importance_score,
            status=entity.status,
            lastAccessedTime=entity.last_accessed_time,
            expiresTime=entity.expires_time,
            createdTime=entity.created_time)


class MemoryListRequestDto(PageRequest):
    scopeType: MemoryScopeType | None = None
    scopeId: str | None = None
    memoryType: str | None = None
    status: MemoryStatus | None = "active"
    ownerAgentId: str | None = None
    userId: str | None = None
    sessionId: str | None = None
    includeExpired: bool = True


class MemoryCreateRequestDto(BaseModel):
    scopeType: MemoryScopeType
    scopeId: str
    memoryType: str = "fact"
    contentText: str
    content: dict[str, JSONValue] | None = None
    metadata: dict[str, JSONValue] = Field(default_factory=dict)
    ownerAgentId: str | None = None
    userId: str | None = None
    sessionId: str | None = None
    sourceRef: str | None = None
    importanceScore: int = Field(default=0, ge=0, le=100)
    expiresTime: datetime | None = None


class MemoryDetailRequestDto(BaseModel):
    memoryId: str


class MemoryUpdateRequestDto(BaseModel):
    memoryId: str
    scopeType: MemoryScopeType | None = None
    scopeId: str | None = None
    memoryType: str | None = None
    contentText: str | None = None
    content: dict[str, JSONValue] | None = None
    metadata: dict[str, JSONValue] | None = None
    ownerAgentId: str | None = None
    userId: str | None = None
    sessionId: str | None = None
    sourceRef: str | None = None
    importanceScore: int | None = Field(default=None, ge=0, le=100)
    expiresTime: datetime | None = None


class MemoryStatusRequestDto(BaseModel):
    memoryId: str
    status: MemoryStatus


class MemoryDeleteRequestDto(BaseModel):
    memoryId: str


class MemorySearchRequestDto(BaseModel):
    query: str
    scopeType: MemoryScopeType | None = None
    scopeId: str | None = None
    ownerAgentId: str | None = None
    userId: str | None = None
    sessionId: str | None = None
    memoryType: str | None = None
    limit: int = Field(default=8, ge=1, le=100)


class MemorySearchResultDto(BaseModel):
    item: MemoryItemDto
    score: float
    scoreDetails: dict[str, JSONValue] = Field(default_factory=dict)


class DeleteData(BaseModel):
    deleted: bool
    memoryId: str