memory-gateway/memory_gateway/schemas_v2.py

"""Schemas for the Memory Gateway v2 control-plane API."""
from __future__ import annotations

from datetime import datetime
from enum import Enum
from typing import Any, Literal, Optional
from uuid import uuid4

from pydantic import BaseModel, Field

from .schemas import utc_now


class OperationStatus(str, Enum):
    ACCEPTED = "accepted"
    RUNNING = "running"
    SUCCESS = "success"
    PARTIAL_SUCCESS = "partial_success"
    FAILED = "failed"
    PENDING = "pending"
    SKIPPED = "skipped"


class BackendRefStatus(str, Enum):
    PENDING = "pending"
    SUCCESS = "success"
    FAILED = "failed"
    SKIPPED = "skipped"


class BackendType(str, Enum):
    OPENVIKING = "openviking"
    EVERMEMOS = "evermemos"
    OBSIDIAN = "obsidian"


class MemoryRefType(str, Enum):
    SESSION_ARCHIVE = "session_archive"
    CONTEXT_RESOURCE = "context_resource"
    MESSAGE_MEMORY = "message_memory"
    EPISODIC_MEMORY = "episodic_memory"
    PROFILE = "profile"
    LONG_TERM_MEMORY = "long_term_memory"
    DRAFT_REVIEW = "draft_review"


class TraceContext(BaseModel):
    trace_id: Optional[str] = None
    span_id: Optional[str] = None
    parent_span_id: Optional[str] = None
    request_id: Optional[str] = None
    metadata: dict[str, Any] = Field(default_factory=dict)


class IngestPolicy(BaseModel):
    allow_openviking: bool = True
    allow_evermemos: bool = True
    allow_obsidian_review: bool = False
    redact_sensitive: bool = True
    require_human_review: bool = False
    metadata: dict[str, Any] = Field(default_factory=dict)


class IngestRequest(BaseModel):
    workspace_id: str
    user_id: str
    agent_id: str
    session_id: str
    turn_id: str
    request_id: Optional[str] = None
    idempotency_key: Optional[str] = None
    namespace: str
    source_type: str = "conversation"
    source_event_id: Optional[str] = None
    role: Literal["system", "user", "assistant", "tool", "agent"] = "user"
    content: str
    policy: IngestPolicy = Field(default_factory=IngestPolicy)
    trace: TraceContext = Field(default_factory=TraceContext)
    metadata: dict[str, Any] = Field(default_factory=dict)


class MemoryRef(BaseModel):
    id: str = Field(default_factory=lambda: f"ref_{uuid4().hex[:16]}")
    gateway_id: str
    workspace_id: str
    user_id: str
    agent_id: Optional[str] = None
    session_id: Optional[str] = None
    turn_id: Optional[str] = None
    namespace: Optional[str] = None
    backend_type: BackendType
    ref_type: MemoryRefType
    native_id: Optional[str] = None
    native_uri: Optional[str] = None
    provenance_id: Optional[str] = None
    idempotency_key: Optional[str] = None
    content_hash: Optional[str] = None
    source_type: Optional[str] = None
    source_event_id: Optional[str] = None
    status: BackendRefStatus = BackendRefStatus.PENDING
    error_message: Optional[str] = None
    metadata: dict[str, Any] = Field(default_factory=dict)
    created_at: datetime = Field(default_factory=utc_now)
    updated_at: datetime = Field(default_factory=utc_now)


class MemoryRefView(MemoryRef):
    pass


class IngestResponse(BaseModel):
    status: OperationStatus
    gateway_id: str
    provenance_id: str
    request_id: Optional[str] = None
    turn_id: str
    refs: list[MemoryRefView] = Field(default_factory=list)
    errors: list[str] = Field(default_factory=list)
    metadata: dict[str, Any] = Field(default_factory=dict)


class CommitRequest(BaseModel):
    workspace_id: str
    user_id: str
    agent_id: Optional[str] = None
    namespace: Optional[str] = None
    request_id: Optional[str] = None
    idempotency_key: Optional[str] = None
    policy: IngestPolicy = Field(default_factory=IngestPolicy)
    metadata: dict[str, Any] = Field(default_factory=dict)


class CommitResponse(BaseModel):
    status: OperationStatus = OperationStatus.ACCEPTED
    job_id: str
    session_id: str
    message: str = "commit accepted"
    refs: list[MemoryRefView] = Field(default_factory=list)
    metadata: dict[str, Any] = Field(default_factory=dict)


class OutboxSummary(BaseModel):
    total_events: int = 0
    pending_events: int = 0
    processing_events: int = 0
    success_events: int = 0
    skipped_events: int = 0
    dead_letter_events: int = 0


class CommitJobView(BaseModel):
    job_id: str
    workspace_id: str
    user_id: str
    agent_id: Optional[str] = None
    session_id: str
    namespace: Optional[str] = None
    status: OperationStatus
    created_refs_count: int = 0
    error_message: Optional[str] = None
    created_at: datetime
    updated_at: datetime
    started_at: Optional[datetime] = None
    finished_at: Optional[datetime] = None
    outbox_summary: OutboxSummary = Field(default_factory=OutboxSummary)


class OutboxProcessResponse(BaseModel):
    status: OperationStatus
    worker_id: str
    processed_count: int = 0
    outbox_summary: OutboxSummary = Field(default_factory=OutboxSummary)


class RetrieveRequest(BaseModel):
    workspace_id: str
    user_id: str
    agent_id: Optional[str] = None
    session_id: Optional[str] = None
    namespace: Optional[str] = None
    query: str
    limit: int = Field(default=10, ge=1, le=100)
    metadata: dict[str, Any] = Field(default_factory=dict)


class ContextItem(BaseModel):
    text: Optional[str] = None
    source_backend: BackendType
    ref_id: Optional[str] = None
    score: float = 0.0
    memory_type: Optional[str] = None
    metadata: dict[str, Any] = Field(default_factory=dict)


class ContextConflict(BaseModel):
    ref_ids: list[str] = Field(default_factory=list)
    reason: str
    metadata: dict[str, Any] = Field(default_factory=dict)


class RetrieveResponse(BaseModel):
    status: OperationStatus
    items: list[ContextItem] = Field(default_factory=list)
    refs: list[MemoryRefView] = Field(default_factory=list)
    conflicts: list[ContextConflict] = Field(default_factory=list)
    trace_id: Optional[str] = None
    metadata: dict[str, Any] = Field(default_factory=dict)


class FeedbackRequest(BaseModel):
    workspace_id: str
    user_id: str
    agent_id: Optional[str] = None
    session_id: Optional[str] = None
    namespace: Optional[str] = None
    memory_ref_id: Optional[str] = None
    feedback_type: Literal["useful", "not_useful", "incorrect", "duplicate", "outdated", "review_approved", "review_rejected"]
    comment: Optional[str] = None
    source_type: str = "manual"
    source_event_id: Optional[str] = None
    metadata: dict[str, Any] = Field(default_factory=dict)


class FeedbackResponse(BaseModel):
    status: OperationStatus
    feedback_id: str
    memory_ref_id: Optional[str] = None
    metadata: dict[str, Any] = Field(default_factory=dict)