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

from datetime import datetime
from typing import TYPE_CHECKING

from pydantic import BaseModel, Field

from core_domain import (
    MemoryItemContract,
    MemoryScopeType,
    MemorySearchResultContract,
    MemoryStatus,
)
from core_shared import JSONValue

if TYPE_CHECKING:
    from app.db.models import MemoryItem


class MemoryCreateRequest(BaseModel):
    tenant_id: str
    scope_type: MemoryScopeType
    scope_id: str
    memory_type: str = "fact"
    content_text: str
    content_json: dict[str, JSONValue] | None = None
    metadata_json: dict[str, JSONValue] = Field(default_factory=dict)
    owner_agent_id: str | None = None
    user_id: str | None = None
    session_id: str | None = None
    source_ref: str | None = None
    importance_score: int = Field(default=0, ge=0, le=100)
    expires_time: datetime | None = None


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


class MemorySearchRequest(BaseModel):
    tenant_id: str
    query: str
    scope_type: MemoryScopeType | None = None
    scope_id: str | None = None
    owner_agent_id: str | None = None
    user_id: str | None = None
    session_id: str | None = None
    limit: int = Field(default=8, ge=1, le=100)


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) -> "MemorySearchResultResponse":
        return cls(item=MemoryResponse.from_entity(entity), score=score)