Files
resolutionflow/backend/app/models/tree_embedding.py
Michael Chihlas 1aa60dada2 feat: add AI assistant with in-session copilot and standalone chat with RAG
Implements three-phase AI assistant feature:
- Phase 0: RAG infrastructure with pgvector embeddings, Voyage AI integration,
  tree chunking service, and semantic search over team's flow library
- Phase 1: In-session copilot panel during flow navigation with contextual
  AI help, current step awareness, and suggested related flows
- Phase 2: Standalone AI chat page with persistent conversation history,
  pin/delete, and configurable retention policies (account-level)

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-04 01:36:36 -05:00

73 lines
2.4 KiB
Python

"""Tree embedding storage for RAG-powered AI assistant.
Stores vector embeddings of tree content chunks for semantic search.
Each tree is split into multiple chunks (node, solution, tree_summary)
and embedded via Voyage AI for cosine similarity retrieval.
"""
import uuid
from datetime import datetime, timezone
from typing import Optional, Any
from sqlalchemy import String, Text, DateTime, ForeignKey, Index
from sqlalchemy.orm import Mapped, mapped_column
from sqlalchemy.dialects.postgresql import UUID, JSONB
from app.core.database import Base
# pgvector column type — imported at runtime to avoid import errors
# when pgvector is not installed locally
try:
from pgvector.sqlalchemy import Vector
except ImportError:
Vector = None
class TreeEmbedding(Base):
__tablename__ = "tree_embeddings"
__table_args__ = (
Index("ix_tree_embeddings_account_id", "account_id"),
Index("ix_tree_embeddings_tree_id", "tree_id"),
)
id: Mapped[uuid.UUID] = mapped_column(
UUID(as_uuid=True), primary_key=True, default=uuid.uuid4
)
tree_id: Mapped[uuid.UUID] = mapped_column(
UUID(as_uuid=True),
ForeignKey("trees.id", ondelete="CASCADE"),
nullable=False,
)
account_id: Mapped[Optional[uuid.UUID]] = mapped_column(
UUID(as_uuid=True),
ForeignKey("accounts.id", ondelete="CASCADE"),
nullable=True,
)
chunk_type: Mapped[str] = mapped_column(
String(30),
nullable=False,
comment="node | solution | tree_summary",
)
node_type: Mapped[Optional[str]] = mapped_column(
String(30), nullable=True
)
node_id: Mapped[Optional[str]] = mapped_column(
String(100), nullable=True
)
chunk_text: Mapped[str] = mapped_column(Text, nullable=False)
embedding_model: Mapped[str] = mapped_column(
String(50), nullable=False, default="voyage-3.5"
)
# The embedding column is created via migration with vector(1024) type
# We store it as a generic column here and handle it in queries
meta: Mapped[dict[str, Any]] = mapped_column(
JSONB, nullable=False, default=dict
)
created_at: Mapped[datetime] = mapped_column(
DateTime(timezone=True), default=lambda: datetime.now(timezone.utc)
)
updated_at: Mapped[datetime] = mapped_column(
DateTime(timezone=True),
default=lambda: datetime.now(timezone.utc),
onupdate=lambda: datetime.now(timezone.utc),
)