fleet-memory/hindsight-api-slim/hindsight_api/models.py
Nicolò Boschi 15ea23d5d6
feat: introduce hindsight-api-slim and hindsight-all-slim packages (#560)
* feat: introduce hindsight-api-slim and hindsight-all-slim packages

Closes #552

- Move all source code from hindsight-api/ to new hindsight-api-slim/
- hindsight-api-slim has heavy ML deps (torch, sentence-transformers,
  transformers, einops, flashrank, mlx, mlx-lm, safetensors) and
  pg0-embedded as optional extras: [local-ml], [embedded-db], [all]
- hindsight-api becomes a zero-code meta-package depending on
  hindsight-api-slim[all] for full backward compatibility
- Add hindsight-all-slim meta-package: hindsight-api-slim + client + embed
- hindsight-all updated to depend on hindsight-api-slim[all]
- pg0.py: lazy-import pg0 with clear ImportError pointing to [embedded-db]
- Dockerfile: replace sed hack with proper uv sync --extra flags
- Update release.yml, test.yml, lint.sh, release.sh, CLAUDE.md and
  all path references throughout the repo

* refactor: rename hindsight/ directory to hindsight-all/

* docs: document hindsight-api-slim and hindsight-all-slim package variants

Add package variants table and extras explanation to installation.md

* docs: remove emojis from installation.md, use professional tone

* docs: link Docker slim variant to pip package variants section

* docs: consolidate Docker image variants into single table

* ci: fix working-directory paths after package restructure

- Replace all hindsight-api → hindsight-api-slim in test.yml
- Replace hindsight → hindsight-all in test.yml
- Add --extra embedded-db to test-embed API install step

* ci: add local-ml and embedded-db extras to API sync steps

These extras were previously implicit in the old hindsight-api package
(which bundled everything). Now that hindsight-api-slim uses optional
extras, we must explicitly request local-ml and embedded-db in CI.

* ci: add API install step with embedded-db to test-embed smoke test

The smoke test starts hindsight-api as a daemon, which requires pg0-embedded.
Add a dedicated install step for hindsight-api-slim with embedded-db extra
so the daemon can start successfully.

* ci: remove --no-install-project when using optional extras

When --no-install-project is combined with --extra, the optional deps
are not installed because extras require the project to be active.
Remove --no-install-project from steps that need local-ml or embedded-db.

* ci: fix ordering of uv sync steps to preserve optional extras

When uv sync runs for a different workspace member, it removes optional
extras installed for other members. Fix by always running extra-requiring
API sync last, after other workspace member syncs.

Also remove --no-install-project from embedded-db sync in test-embed,
as --no-install-project prevents optional extras from being active.

* ci: add local-ml extra to test-embed API install for smoke test

The smoke test starts the full API server which needs sentence-transformers
for local embeddings (default provider). Add local-ml extra to the install.

* ci: simplify extras with --all-extras and add slim pip smoke test

- Replace explicit --extra local-ml --extra embedded-db with --all-extras
  for cleaner, more maintainable sync steps
- Add test-pip-slim job: tests hindsight-api-slim[embedded-db] without
  local ML models, using Cohere for embeddings/reranking (mirrors Docker
  slim smoke test approach)

* ci: simplify slim smoke test to health check only (mirrors Docker test)
2026-03-13 13:50:03 +01:00

316 lines
13 KiB
Python

"""
SQLAlchemy models for the memory system.
"""
from dataclasses import dataclass
from datetime import datetime
from uuid import UUID as PyUUID
@dataclass
class RequestContext:
"""
Context for request authentication and authorization.
This dataclass carries authentication data from HTTP requests to the
memory engine operations. It can be extended to include additional
context like headers, tokens, user info, etc.
"""
api_key: str | None = None
api_key_id: str | None = None # UUID of the API key used for authentication
tenant_id: str | None = None # Tenant identifier (set by extension after auth)
internal: bool = False # True for background/internal operations (skips extension auth)
user_initiated: bool = False # True for async operations that originated from a user request
allowed_bank_ids: list[str] | None = None # None = unrestricted (all banks)
from pgvector.sqlalchemy import Vector
from sqlalchemy import (
CheckConstraint,
Float,
ForeignKey,
ForeignKeyConstraint,
Index,
Integer,
Text,
func,
)
from sqlalchemy import (
text as sql_text,
)
from sqlalchemy.dialects.postgresql import JSONB, TIMESTAMP, UUID
from sqlalchemy.ext.asyncio import AsyncAttrs
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, relationship
from .config import EMBEDDING_DIMENSION
class Base(AsyncAttrs, DeclarativeBase):
"""Base class for all models."""
pass
class Document(Base):
"""Source documents for memory units."""
__tablename__ = "documents"
id: Mapped[str] = mapped_column(Text, primary_key=True)
bank_id: Mapped[str] = mapped_column(Text, primary_key=True)
original_text: Mapped[str | None] = mapped_column(Text)
content_hash: Mapped[str | None] = mapped_column(Text)
doc_metadata: Mapped[dict] = mapped_column("metadata", JSONB, server_default=sql_text("'{}'::jsonb"))
created_at: Mapped[datetime] = mapped_column(TIMESTAMP(timezone=True), server_default=func.now())
updated_at: Mapped[datetime] = mapped_column(TIMESTAMP(timezone=True), server_default=func.now())
# Relationships
memory_units = relationship("MemoryUnit", back_populates="document", cascade="all, delete-orphan")
__table_args__ = (
Index("idx_documents_bank_id", "bank_id"),
Index("idx_documents_content_hash", "content_hash"),
)
class MemoryUnit(Base):
"""Individual sentence-level memories."""
__tablename__ = "memory_units"
id: Mapped[PyUUID] = mapped_column(
UUID(as_uuid=True), primary_key=True, server_default=sql_text("gen_random_uuid()")
)
bank_id: Mapped[str] = mapped_column(Text, nullable=False)
document_id: Mapped[str | None] = mapped_column(Text)
text: Mapped[str] = mapped_column(Text, nullable=False)
embedding = mapped_column(Vector(EMBEDDING_DIMENSION)) # pgvector type
context: Mapped[str | None] = mapped_column(Text)
event_date: Mapped[datetime] = mapped_column(
TIMESTAMP(timezone=True), nullable=False
) # Kept for backward compatibility
occurred_start: Mapped[datetime | None] = mapped_column(
TIMESTAMP(timezone=True)
) # When fact occurred (range start)
occurred_end: Mapped[datetime | None] = mapped_column(TIMESTAMP(timezone=True)) # When fact occurred (range end)
mentioned_at: Mapped[datetime | None] = mapped_column(TIMESTAMP(timezone=True)) # When fact was mentioned
fact_type: Mapped[str] = mapped_column(Text, nullable=False, server_default="world")
confidence_score: Mapped[float | None] = mapped_column(Float)
unit_metadata: Mapped[dict] = mapped_column(
"metadata", JSONB, server_default=sql_text("'{}'::jsonb")
) # User-defined metadata (str->str)
created_at: Mapped[datetime] = mapped_column(TIMESTAMP(timezone=True), server_default=func.now())
updated_at: Mapped[datetime] = mapped_column(TIMESTAMP(timezone=True), server_default=func.now())
# Relationships
document = relationship("Document", back_populates="memory_units")
unit_entities = relationship("UnitEntity", back_populates="memory_unit", cascade="all, delete-orphan")
outgoing_links = relationship(
"MemoryLink", foreign_keys="MemoryLink.from_unit_id", back_populates="from_unit", cascade="all, delete-orphan"
)
incoming_links = relationship(
"MemoryLink", foreign_keys="MemoryLink.to_unit_id", back_populates="to_unit", cascade="all, delete-orphan"
)
__table_args__ = (
ForeignKeyConstraint(
["document_id", "bank_id"],
["documents.id", "documents.bank_id"],
name="memory_units_document_fkey",
ondelete="CASCADE",
),
CheckConstraint("fact_type IN ('world', 'experience', 'opinion', 'observation')"),
CheckConstraint("confidence_score IS NULL OR (confidence_score >= 0.0 AND confidence_score <= 1.0)"),
CheckConstraint(
"(fact_type = 'opinion' AND confidence_score IS NOT NULL) OR "
"(fact_type = 'observation') OR "
"(fact_type NOT IN ('opinion', 'observation') AND confidence_score IS NULL)",
name="confidence_score_fact_type_check",
),
Index("idx_memory_units_bank_id", "bank_id"),
Index("idx_memory_units_document_id", "document_id"),
Index("idx_memory_units_event_date", "event_date", postgresql_ops={"event_date": "DESC"}),
Index("idx_memory_units_bank_date", "bank_id", "event_date", postgresql_ops={"event_date": "DESC"}),
Index("idx_memory_units_fact_type", "fact_type"),
Index("idx_memory_units_bank_fact_type", "bank_id", "fact_type"),
Index(
"idx_memory_units_bank_type_date",
"bank_id",
"fact_type",
"event_date",
postgresql_ops={"event_date": "DESC"},
),
Index(
"idx_memory_units_opinion_confidence",
"bank_id",
"confidence_score",
postgresql_where=sql_text("fact_type = 'opinion'"),
postgresql_ops={"confidence_score": "DESC"},
),
Index(
"idx_memory_units_opinion_date",
"bank_id",
"event_date",
postgresql_where=sql_text("fact_type = 'opinion'"),
postgresql_ops={"event_date": "DESC"},
),
Index(
"idx_memory_units_observation_date",
"bank_id",
"event_date",
postgresql_where=sql_text("fact_type = 'observation'"),
postgresql_ops={"event_date": "DESC"},
),
Index(
"idx_memory_units_embedding",
"embedding",
postgresql_using="hnsw",
postgresql_ops={"embedding": "vector_cosine_ops"},
),
)
class Entity(Base):
"""Resolved entities (people, organizations, locations, etc.)."""
__tablename__ = "entities"
id: Mapped[PyUUID] = mapped_column(
UUID(as_uuid=True), primary_key=True, server_default=sql_text("gen_random_uuid()")
)
canonical_name: Mapped[str] = mapped_column(Text, nullable=False)
bank_id: Mapped[str] = mapped_column(Text, nullable=False)
entity_metadata: Mapped[dict] = mapped_column("metadata", JSONB, server_default=sql_text("'{}'::jsonb"))
first_seen: Mapped[datetime] = mapped_column(TIMESTAMP(timezone=True), server_default=func.now())
last_seen: Mapped[datetime] = mapped_column(TIMESTAMP(timezone=True), server_default=func.now())
mention_count: Mapped[int] = mapped_column(Integer, server_default="1")
# Relationships
unit_entities = relationship("UnitEntity", back_populates="entity", cascade="all, delete-orphan")
memory_links = relationship("MemoryLink", back_populates="entity", cascade="all, delete-orphan")
cooccurrences_1 = relationship(
"EntityCooccurrence",
foreign_keys="EntityCooccurrence.entity_id_1",
back_populates="entity_1",
cascade="all, delete-orphan",
)
cooccurrences_2 = relationship(
"EntityCooccurrence",
foreign_keys="EntityCooccurrence.entity_id_2",
back_populates="entity_2",
cascade="all, delete-orphan",
)
__table_args__ = (
Index("idx_entities_bank_id", "bank_id"),
Index("idx_entities_canonical_name", "canonical_name"),
Index("idx_entities_bank_name", "bank_id", "canonical_name"),
)
class UnitEntity(Base):
"""Association between memory units and entities."""
__tablename__ = "unit_entities"
unit_id: Mapped[PyUUID] = mapped_column(
UUID(as_uuid=True), ForeignKey("memory_units.id", ondelete="CASCADE"), primary_key=True
)
entity_id: Mapped[PyUUID] = mapped_column(
UUID(as_uuid=True), ForeignKey("entities.id", ondelete="CASCADE"), primary_key=True
)
# Relationships
memory_unit = relationship("MemoryUnit", back_populates="unit_entities")
entity = relationship("Entity", back_populates="unit_entities")
__table_args__ = (
Index("idx_unit_entities_unit", "unit_id"),
Index("idx_unit_entities_entity", "entity_id"),
)
class EntityCooccurrence(Base):
"""Materialized cache of entity co-occurrences."""
__tablename__ = "entity_cooccurrences"
entity_id_1: Mapped[PyUUID] = mapped_column(
UUID(as_uuid=True), ForeignKey("entities.id", ondelete="CASCADE"), primary_key=True
)
entity_id_2: Mapped[PyUUID] = mapped_column(
UUID(as_uuid=True), ForeignKey("entities.id", ondelete="CASCADE"), primary_key=True
)
cooccurrence_count: Mapped[int] = mapped_column(Integer, server_default="1")
last_cooccurred: Mapped[datetime] = mapped_column(TIMESTAMP(timezone=True), server_default=func.now())
# Relationships
entity_1 = relationship("Entity", foreign_keys=[entity_id_1], back_populates="cooccurrences_1")
entity_2 = relationship("Entity", foreign_keys=[entity_id_2], back_populates="cooccurrences_2")
__table_args__ = (
CheckConstraint("entity_id_1 < entity_id_2", name="entity_cooccurrence_order_check"),
Index("idx_entity_cooccurrences_entity1", "entity_id_1"),
Index("idx_entity_cooccurrences_entity2", "entity_id_2"),
Index("idx_entity_cooccurrences_count", "cooccurrence_count", postgresql_ops={"cooccurrence_count": "DESC"}),
)
class MemoryLink(Base):
"""Links between memory units (temporal, semantic, entity)."""
__tablename__ = "memory_links"
from_unit_id: Mapped[PyUUID] = mapped_column(
UUID(as_uuid=True), ForeignKey("memory_units.id", ondelete="CASCADE"), primary_key=True
)
to_unit_id: Mapped[PyUUID] = mapped_column(
UUID(as_uuid=True), ForeignKey("memory_units.id", ondelete="CASCADE"), primary_key=True
)
link_type: Mapped[str] = mapped_column(Text, primary_key=True)
entity_id: Mapped[PyUUID | None] = mapped_column(
UUID(as_uuid=True), ForeignKey("entities.id", ondelete="CASCADE"), primary_key=True
)
weight: Mapped[float] = mapped_column(Float, nullable=False, server_default="1.0")
created_at: Mapped[datetime] = mapped_column(TIMESTAMP(timezone=True), server_default=func.now())
# Relationships
from_unit = relationship("MemoryUnit", foreign_keys=[from_unit_id], back_populates="outgoing_links")
to_unit = relationship("MemoryUnit", foreign_keys=[to_unit_id], back_populates="incoming_links")
entity = relationship("Entity", back_populates="memory_links")
__table_args__ = (
CheckConstraint(
"link_type IN ('temporal', 'semantic', 'entity', 'causes', 'caused_by', 'enables', 'prevents')",
name="memory_links_link_type_check",
),
CheckConstraint("weight >= 0.0 AND weight <= 1.0", name="memory_links_weight_check"),
Index("idx_memory_links_from", "from_unit_id"),
Index("idx_memory_links_to", "to_unit_id"),
Index("idx_memory_links_type", "link_type"),
Index("idx_memory_links_entity", "entity_id", postgresql_where=sql_text("entity_id IS NOT NULL")),
Index(
"idx_memory_links_from_weight",
"from_unit_id",
"weight",
postgresql_where=sql_text("weight >= 0.1"),
postgresql_ops={"weight": "DESC"},
),
)
class Bank(Base):
"""Memory bank profiles with disposition traits and background."""
__tablename__ = "banks"
bank_id: Mapped[str] = mapped_column(Text, primary_key=True)
disposition: Mapped[dict] = mapped_column(
JSONB, nullable=False, server_default=sql_text('\'{"skepticism": 3, "literalism": 3, "empathy": 3}\'::jsonb')
)
background: Mapped[str] = mapped_column(Text, nullable=False, server_default="")
created_at: Mapped[datetime] = mapped_column(TIMESTAMP(timezone=True), server_default=func.now())
updated_at: Mapped[datetime] = mapped_column(TIMESTAMP(timezone=True), server_default=func.now())
__table_args__ = (Index("idx_banks_bank_id", "bank_id"),)