blackroad-operating-system/backend/app/models/cognition.py

"""
Cognition Database Models

Models for storing:
- Workflows and their execution history
- Reasoning traces from agents
- Agent memory/context
- Prompt registry

Tables:
- workflows: Workflow definitions and execution status
- workflow_executions: History of workflow runs
- reasoning_traces: Agent reasoning step records
- agent_memory: Shared memory/context across workflow
- prompt_registry: Registered agent prompts
"""

from datetime import datetime
from sqlalchemy import Column, Integer, String, Text, Float, Boolean, DateTime, ForeignKey, JSON, Enum
from sqlalchemy.orm import relationship
from sqlalchemy.dialects.postgresql import UUID, JSONB
import uuid
import enum

from ..database import Base


class WorkflowStatus(str, enum.Enum):
    """Workflow execution status"""
    PENDING = "pending"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"
    CANCELLED = "cancelled"


class ExecutionMode(str, enum.Enum):
    """Workflow execution mode"""
    SEQUENTIAL = "sequential"
    PARALLEL = "parallel"
    RECURSIVE = "recursive"


class Workflow(Base):
    """
    Workflow Definition

    Stores multi-agent workflow definitions.
    """
    __tablename__ = "workflows"

    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    name = Column(String(200), nullable=False, index=True)
    description = Column(Text)

    # Workflow configuration
    mode = Column(Enum(ExecutionMode), default=ExecutionMode.SEQUENTIAL, nullable=False)
    steps = Column(JSONB, nullable=False)  # List of workflow steps
    timeout_seconds = Column(Integer, default=600)

    # Metadata
    created_by = Column(String(100))
    created_at = Column(DateTime, default=datetime.utcnow, nullable=False)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
    is_active = Column(Boolean, default=True)
    is_template = Column(Boolean, default=False)

    # Tags for categorization
    tags = Column(JSONB, default=list)

    # Relationships
    executions = relationship("WorkflowExecution", back_populates="workflow", cascade="all, delete-orphan")

    def __repr__(self):
        return f"<Workflow(id={self.id}, name={self.name}, mode={self.mode})>"


class WorkflowExecution(Base):
    """
    Workflow Execution Record

    Stores history of workflow executions with results.
    """
    __tablename__ = "workflow_executions"

    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    workflow_id = Column(UUID(as_uuid=True), ForeignKey("workflows.id"), nullable=False, index=True)

    # Execution details
    status = Column(Enum(WorkflowStatus), default=WorkflowStatus.PENDING, nullable=False, index=True)
    started_at = Column(DateTime, default=datetime.utcnow, nullable=False, index=True)
    completed_at = Column(DateTime)
    duration_seconds = Column(Float)

    # Results
    step_results = Column(JSONB)  # Results from each step
    error_message = Column(Text)
    error_details = Column(JSONB)

    # Metrics
    overall_confidence = Column(Float)
    total_agents_used = Column(Integer)

    # Context
    initial_context = Column(JSONB)
    final_memory = Column(JSONB)

    # Relationships
    workflow = relationship("Workflow", back_populates="executions")
    reasoning_traces = relationship("ReasoningTrace", back_populates="execution", cascade="all, delete-orphan")

    def __repr__(self):
        return f"<WorkflowExecution(id={self.id}, workflow_id={self.workflow_id}, status={self.status})>"


class ReasoningTrace(Base):
    """
    Reasoning Trace Step

    Stores individual reasoning steps from agent execution.
    Provides transparency into how agents arrived at decisions.
    """
    __tablename__ = "reasoning_traces"

    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    execution_id = Column(UUID(as_uuid=True), ForeignKey("workflow_executions.id"), nullable=False, index=True)

    # Step identification
    workflow_step_name = Column(String(100), nullable=False)
    agent_name = Column(String(50), nullable=False, index=True)
    step_number = Column(Integer, nullable=False)
    step_name = Column(String(100), nullable=False)
    step_emoji = Column(String(10))

    # Reasoning data
    input_context = Column(Text)
    output = Column(Text)
    confidence_score = Column(Float)

    # Additional metadata
    metadata = Column(JSONB)
    timestamp = Column(DateTime, default=datetime.utcnow, nullable=False, index=True)

    # Relationships
    execution = relationship("WorkflowExecution", back_populates="reasoning_traces")

    def __repr__(self):
        return f"<ReasoningTrace(id={self.id}, agent={self.agent_name}, step={self.step_name})>"


class AgentMemory(Base):
    """
    Agent Memory/Context

    Stores shared context and memory across workflow execution.
    Enables agents to build upon each other's work.
    """
    __tablename__ = "agent_memory"

    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    execution_id = Column(UUID(as_uuid=True), ForeignKey("workflow_executions.id"), index=True)

    # Memory data
    context = Column(JSONB, nullable=False)  # Shared context dictionary
    confidence_scores = Column(JSONB)  # Confidence per step

    # Metadata
    created_at = Column(DateTime, default=datetime.utcnow, nullable=False)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)

    # Memory can be associated with user sessions
    session_id = Column(String(100), index=True)
    user_id = Column(String(100), index=True)

    # TTL for memory expiration
    expires_at = Column(DateTime)

    def __repr__(self):
        return f"<AgentMemory(id={self.id}, execution_id={self.execution_id})>"


class PromptRegistry(Base):
    """
    Prompt Registry

    Stores registered agent prompts (summon spells).
    Enables versioning and management of agent invocation prompts.
    """
    __tablename__ = "prompt_registry"

    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)

    # Prompt identification
    agent_name = Column(String(50), nullable=False, index=True)
    prompt_name = Column(String(100))
    prompt_text = Column(Text, nullable=False)

    # Versioning
    version = Column(String(20), nullable=False)
    is_active = Column(Boolean, default=True, index=True)

    # Metadata
    description = Column(Text)
    metadata = Column(JSONB)  # Author, purpose, etc.
    tags = Column(JSONB, default=list)

    # Usage stats
    usage_count = Column(Integer, default=0)
    last_used_at = Column(DateTime)
    average_confidence = Column(Float)

    # Timestamps
    created_at = Column(DateTime, default=datetime.utcnow, nullable=False, index=True)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
    created_by = Column(String(100))

    def __repr__(self):
        return f"<PromptRegistry(id={self.id}, agent={self.agent_name}, version={self.version})>"


class AgentPerformanceMetric(Base):
    """
    Agent Performance Metrics

    Tracks performance metrics for agents over time.
    Enables monitoring and optimization.
    """
    __tablename__ = "agent_performance_metrics"

    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)

    # Agent identification
    agent_name = Column(String(50), nullable=False, index=True)
    execution_id = Column(UUID(as_uuid=True), ForeignKey("workflow_executions.id"), index=True)

    # Performance metrics
    execution_time_seconds = Column(Float)
    confidence_score = Column(Float)
    success = Column(Boolean, default=True)

    # Resource usage (if available)
    memory_usage_mb = Column(Float)
    api_calls_made = Column(Integer)

    # Quality metrics
    reasoning_steps_count = Column(Integer)
    complexity_score = Column(Float)

    # Timestamps
    measured_at = Column(DateTime, default=datetime.utcnow, nullable=False, index=True)

    def __repr__(self):
        return f"<AgentPerformanceMetric(agent={self.agent_name}, confidence={self.confidence_score})>"