""" DAARION Memory Service - Pydantic Models """ from datetime import datetime from typing import Optional, List, Any from uuid import UUID from enum import Enum from pydantic import BaseModel, Field # ============================================================================ # ENUMS # ============================================================================ class EventType(str, Enum): MESSAGE = "message" TOOL_CALL = "tool_call" TOOL_RESULT = "tool_result" DECISION = "decision" SUMMARY = "summary" MEMORY_WRITE = "memory_write" MEMORY_RETRACT = "memory_retract" ERROR = "error" class MessageRole(str, Enum): USER = "user" ASSISTANT = "assistant" SYSTEM = "system" TOOL = "tool" class MemoryCategory(str, Enum): PREFERENCE = "preference" IDENTITY = "identity" CONSTRAINT = "constraint" PROJECT_FACT = "project_fact" RELATIONSHIP = "relationship" SKILL = "skill" GOAL = "goal" CONTEXT = "context" FEEDBACK = "feedback" class RetentionPolicy(str, Enum): PERMANENT = "permanent" SESSION = "session" TTL_DAYS = "ttl_days" UNTIL_REVOKED = "until_revoked" class FeedbackAction(str, Enum): CONFIRM = "confirm" REJECT = "reject" EDIT = "edit" DELETE = "delete" # ============================================================================ # REQUEST MODELS # ============================================================================ class CreateThreadRequest(BaseModel): """Create new conversation thread""" org_id: UUID workspace_id: Optional[UUID] = None user_id: UUID agent_id: Optional[UUID] = None title: Optional[str] = None tags: List[str] = [] metadata: dict = {} class AddEventRequest(BaseModel): """Add event to conversation""" thread_id: UUID event_type: EventType role: Optional[MessageRole] = None content: Optional[str] = None tool_name: Optional[str] = None tool_input: Optional[dict] = None tool_output: Optional[dict] = None payload: dict = {} token_count: Optional[int] = None model_used: Optional[str] = None latency_ms: Optional[int] = None metadata: dict = {} class CreateMemoryRequest(BaseModel): """Create long-term memory item""" org_id: UUID workspace_id: Optional[UUID] = None user_id: UUID agent_id: Optional[UUID] = None # null = global category: MemoryCategory fact_text: str confidence: float = Field(default=0.8, ge=0, le=1) source_event_id: Optional[UUID] = None source_thread_id: Optional[UUID] = None extraction_method: str = "explicit" is_sensitive: bool = False retention: RetentionPolicy = RetentionPolicy.UNTIL_REVOKED ttl_days: Optional[int] = None tags: List[str] = [] metadata: dict = {} class MemoryFeedbackRequest(BaseModel): """User feedback on memory""" memory_id: UUID user_id: UUID action: FeedbackAction new_value: Optional[str] = None reason: Optional[str] = None class RetrievalRequest(BaseModel): """Semantic retrieval request""" org_id: UUID user_id: UUID agent_id: Optional[UUID] = None workspace_id: Optional[UUID] = None queries: List[str] top_k: int = 10 min_confidence: float = 0.5 include_global: bool = True categories: Optional[List[MemoryCategory]] = None class SummaryRequest(BaseModel): """Generate summary for thread""" thread_id: UUID force: bool = False # force even if under token threshold # ============================================================================ # RESPONSE MODELS # ============================================================================ class ThreadResponse(BaseModel): thread_id: UUID org_id: UUID workspace_id: Optional[UUID] user_id: UUID agent_id: Optional[UUID] title: Optional[str] status: str message_count: int total_tokens: int created_at: datetime last_activity_at: datetime class EventResponse(BaseModel): event_id: UUID thread_id: UUID event_type: EventType role: Optional[MessageRole] content: Optional[str] tool_name: Optional[str] tool_input: Optional[dict] tool_output: Optional[dict] payload: dict token_count: Optional[int] sequence_num: int created_at: datetime class MemoryResponse(BaseModel): memory_id: UUID org_id: UUID workspace_id: Optional[UUID] user_id: UUID agent_id: Optional[UUID] category: MemoryCategory fact_text: str confidence: float is_verified: bool is_sensitive: bool retention: RetentionPolicy valid_from: datetime valid_to: Optional[datetime] last_used_at: Optional[datetime] use_count: int created_at: datetime class SummaryResponse(BaseModel): summary_id: UUID thread_id: UUID version: int summary_text: str state: dict events_count: int compression_ratio: Optional[float] created_at: datetime class RetrievalResult(BaseModel): """Single retrieval result""" memory_id: UUID fact_text: str category: MemoryCategory confidence: float relevance_score: float agent_id: Optional[UUID] is_global: bool class RetrievalResponse(BaseModel): """Retrieval response with results""" results: List[RetrievalResult] query_count: int total_results: int class ContextResponse(BaseModel): """Full context for agent prompt""" thread_id: UUID summary: Optional[SummaryResponse] recent_messages: List[EventResponse] retrieved_memories: List[RetrievalResult] token_estimate: int # ============================================================================ # INTERNAL MODELS # ============================================================================ class EmbeddingRequest(BaseModel): """Internal embedding request""" texts: List[str] input_type: str = "search_document" # or "search_query" class QdrantPayload(BaseModel): """Qdrant point payload""" org_id: str workspace_id: Optional[str] user_id: str agent_id: Optional[str] thread_id: Optional[str] memory_id: Optional[str] event_id: Optional[str] type: str # "memory", "summary", "message" category: Optional[str] text: str created_at: str