AI-Writie-Assistant/server/models.py

"""SQLAlchemy models matching the Go project's database schema."""

from datetime import datetime
from sqlalchemy import Column, String, Integer, Boolean, Text, DateTime, create_engine
from sqlalchemy.orm import DeclarativeBase, Session


class Base(DeclarativeBase):
    pass


# ── Global DB Models ──


class LLMProvider(Base):
    __tablename__ = "llm_providers"

    id = Column(String, primary_key=True)
    name = Column(String, default="")
    provider = Column(String, default="")  # Ollama, DeepSeek, OpenAI, Qwen
    base_url = Column(String, default="")
    api_key = Column(String, default="")
    model_id = Column(String, default="")
    enabled = Column(Boolean, default=True)
    created_at = Column(DateTime, default=datetime.utcnow)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)


class VectorDBConfig(Base):
    __tablename__ = "vector_db_configs"

    id = Column(Integer, primary_key=True, default=1)
    endpoint = Column(String, default="")
    api_key = Column(String, default="")
    status = Column(String, default="disconnected")


class EmbeddingModelConfig(Base):
    __tablename__ = "embedding_model_configs"

    id = Column(Integer, primary_key=True, default=1)
    provider = Column(String, default="OpenAI")
    base_url = Column(String, default="")
    api_key = Column(String, default="")
    model_id = Column(String, default="")
    enabled = Column(Boolean, default=False)
    dimensions = Column(Integer, default=1024)


class Project(Base):
    __tablename__ = "projects"

    id = Column(String, primary_key=True)
    name = Column(String, default="")
    path = Column(String, default="")
    created_at = Column(DateTime, default=datetime.utcnow)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)


# ── Project-scoped DB Models ──


class SourceFile(Base):
    __tablename__ = "source_files"

    id = Column(String, primary_key=True)
    project_id = Column(String, index=True)
    name = Column(String, default="")
    type = Column(String, default="")
    category = Column(String, default="")
    file_path = Column(String, default="")
    size = Column(String, default="")
    parsed_content = Column(Text, default="")
    vector_status = Column(String, default="pending")
    created_at = Column(DateTime, default=datetime.utcnow)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)


class ChatMessage(Base):
    __tablename__ = "chat_messages"

    id = Column(Integer, primary_key=True, autoincrement=True)
    project_id = Column(String, index=True)
    role = Column(String, default="")
    content = Column(Text, default="")
    sources = Column(Text, default="")
    citations = Column(Text, default="")
    created_at = Column(DateTime, default=datetime.utcnow)


class TemplateChapter(Base):
    __tablename__ = "template_chapters"

    id = Column(String, primary_key=True)
    project_id = Column(String, index=True)
    template_name = Column(String, default="")
    title = Column(String, default="")
    status = Column(String, default="idle")
    progress = Column(Integer, default=0)
    content = Column(Text, default="")
    sort_order = Column(Integer, default=0)
    created_at = Column(DateTime, default=datetime.utcnow)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)


class TextChunk(Base):
    __tablename__ = "text_chunks"

    id = Column(String, primary_key=True)
    project_id = Column(String, index=True)
    source_id = Column(String, index=True)
    content = Column(Text, default="")
    chunk_idx = Column(Integer, default=0)
    created_at = Column(DateTime, default=datetime.utcnow)


class DeliveryStandard(Base):
    __tablename__ = "delivery_standards"

    id = Column(Integer, primary_key=True, autoincrement=True)
    project_id = Column(String, index=True)
    file_name = Column(String, default="")
    content = Column(Text, default="")
    created_at = Column(DateTime, default=datetime.utcnow)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)


# ── Pydantic schemas for API ──

from pydantic import BaseModel
from typing import Optional


class LLMProviderSchema(BaseModel):
    id: str = ""
    name: str = ""
    provider: str = ""
    url: str = ""
    key: str = ""
    model: str = ""
    enabled: bool = True

    model_config = {"from_attributes": True}

    @classmethod
    def from_orm_obj(cls, obj: LLMProvider) -> "LLMProviderSchema":
        return cls(id=obj.id, name=obj.name, provider=obj.provider,
                   url=obj.base_url, key=obj.api_key, model=obj.model_id,
                   enabled=obj.enabled)


class VectorDBConfigSchema(BaseModel):
    endpoint: str = ""
    apiKey: str = ""
    status: str = "disconnected"


class EmbeddingConfigSchema(BaseModel):
    provider: str = "OpenAI"
    url: str = ""
    key: str = ""
    model: str = ""
    enabled: bool = False
    dimensions: int = 1024

    @classmethod
    def from_orm_obj(cls, obj: EmbeddingModelConfig) -> "EmbeddingConfigSchema":
        return cls(provider=obj.provider, url=obj.base_url, key=obj.api_key,
                   model=obj.model_id, enabled=obj.enabled,
                   dimensions=obj.dimensions or 1024)


class ProjectSchema(BaseModel):
    id: str
    name: str

    model_config = {"from_attributes": True}


class SourceFileSchema(BaseModel):
    id: str
    name: str
    type: str
    category: str
    size: str
    vectorStatus: str = "pending"

    model_config = {"from_attributes": True}

    @classmethod
    def from_orm_obj(cls, obj: SourceFile) -> "SourceFileSchema":
        return cls(id=obj.id, name=obj.name, type=obj.type,
                   category=obj.category, size=obj.size,
                   vectorStatus=obj.vector_status)


class ChatMessageSchema(BaseModel):
    id: int = 0
    projectId: str = ""
    role: str = ""
    content: str = ""
    sources: str = ""
    citations: str = ""

    model_config = {"from_attributes": True}

    @classmethod
    def from_orm_obj(cls, obj: ChatMessage) -> "ChatMessageSchema":
        return cls(id=obj.id, projectId=obj.project_id, role=obj.role,
                   content=obj.content, sources=obj.sources or "",
                   citations=obj.citations or "")


class TemplateChapterSchema(BaseModel):
    id: str
    title: str
    status: str = "idle"
    progress: int = 0
    content: str = ""
    sortOrder: int = 0

    model_config = {"from_attributes": True}

    @classmethod
    def from_orm_obj(cls, obj: TemplateChapter) -> "TemplateChapterSchema":
        return cls(id=obj.id, title=obj.title, status=obj.status,
                   progress=obj.progress, content=obj.content,
                   sortOrder=obj.sort_order)


class DeliveryStandardSchema(BaseModel):
    fileName: str = ""
    content: str = ""