From 70de23ed8566e57797fea17d0c62252d8de732a7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Boschi?= Date: Mon, 5 Jan 2026 14:43:05 +0100 Subject: [PATCH] feat: configurable embedding dimensions + OpenAI Embeddings (#101) * feat: configurable embedding dimensions + OpenAI Embeddings * fix tests --- .github/ISSUE_TEMPLATE/bug_report.yml | 71 +++ .github/ISSUE_TEMPLATE/config.yml | 8 + .github/ISSUE_TEMPLATE/feature_request.yml | 82 ++++ .github/workflows/test.yml | 1 + hindsight-api/hindsight_api/config.py | 8 +- .../hindsight_api/engine/embeddings.py | 202 +++++++-- .../hindsight_api/engine/memory_engine.py | 6 +- hindsight-api/hindsight_api/migrations.py | 128 ++++++ hindsight-api/hindsight_api/models.py | 4 +- .../tests/test_custom_embedding_dimension.py | 428 ++++++++++++++++++ .../docs/developer/configuration.md | 25 +- uv.lock | 10 +- 12 files changed, 935 insertions(+), 38 deletions(-) create mode 100644 .github/ISSUE_TEMPLATE/bug_report.yml create mode 100644 .github/ISSUE_TEMPLATE/config.yml create mode 100644 .github/ISSUE_TEMPLATE/feature_request.yml create mode 100644 hindsight-api/tests/test_custom_embedding_dimension.py diff --git a/.github/ISSUE_TEMPLATE/bug_report.yml b/.github/ISSUE_TEMPLATE/bug_report.yml new file mode 100644 index 00000000..40a5bf10 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/bug_report.yml @@ -0,0 +1,71 @@ +name: Bug Report +description: Report a bug or unexpected behavior +labels: ["bug", "triage"] +body: + - type: markdown + attributes: + value: | + Thanks for taking the time to report a bug! Please fill out the sections below. + + - type: textarea + id: description + attributes: + label: Bug Description + description: A clear and concise description of the bug + placeholder: What happened? + validations: + required: true + + - type: textarea + id: reproduction + attributes: + label: Steps to Reproduce + description: Steps to reproduce the behavior + placeholder: | + 1. Configure '...' + 2. Call '...' + 3. See error + validations: + required: true + + - type: textarea + id: expected + attributes: + label: Expected Behavior + description: What did you expect to happen? + validations: + required: true + + - type: textarea + id: actual + attributes: + label: Actual Behavior + description: What actually happened? + validations: + required: true + + - type: input + id: version + attributes: + label: Version + description: What version are you using? + placeholder: e.g., 0.1.0 or commit hash + validations: + required: false + + - type: dropdown + id: llm-provider + attributes: + label: LLM Provider + description: Which LLM provider are you using? + options: + - OpenAI + - Anthropic + - Gemini + - Groq + - Ollama + - LM Studio + - Other + validations: + required: false + diff --git a/.github/ISSUE_TEMPLATE/config.yml b/.github/ISSUE_TEMPLATE/config.yml new file mode 100644 index 00000000..e79e28e6 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/config.yml @@ -0,0 +1,8 @@ +blank_issues_enabled: false +contact_links: + - name: Questions & Help + url: https://github.com/vectorize-io/hindsight/discussions/categories/q-a + about: Please ask questions and get help in Discussions instead of opening an issue. + - name: Ideas & Feedback + url: https://github.com/vectorize-io/hindsight/discussions/categories/ideas + about: Share ideas or give feedback in Discussions. diff --git a/.github/ISSUE_TEMPLATE/feature_request.yml b/.github/ISSUE_TEMPLATE/feature_request.yml new file mode 100644 index 00000000..b5c9faaf --- /dev/null +++ b/.github/ISSUE_TEMPLATE/feature_request.yml @@ -0,0 +1,82 @@ +name: Feature Request +description: Suggest a new feature or enhancement +labels: ["enhancement", "triage"] +body: + - type: markdown + attributes: + value: | + Thanks for suggesting a feature! Please describe what you'd like to see added. + + - type: textarea + id: use-case + attributes: + label: Use Case + description: Describe your specific use case. What are you building? What's your goal? + placeholder: | + I'm building an AI agent that needs to... + My application handles... + validations: + required: true + + - type: textarea + id: problem + attributes: + label: Problem Statement + description: What problem are you facing? What's missing or difficult today? + placeholder: Currently I have to... which causes... + validations: + required: true + + - type: textarea + id: benefit + attributes: + label: How This Feature Would Help + description: Explain how this feature would improve your workflow or solve your problem + placeholder: With this feature, I would be able to... + validations: + required: true + + - type: textarea + id: solution + attributes: + label: Proposed Solution + description: Describe your ideal solution (optional - we may have ideas too!) + placeholder: It would be great if Hindsight could... + validations: + required: false + + - type: textarea + id: alternatives + attributes: + label: Alternatives Considered + description: Have you considered any alternative solutions or workarounds? + validations: + required: false + + - type: dropdown + id: priority + attributes: + label: Priority + description: How important is this feature to you? + options: + - Nice to have + - Important - affects my workflow + - Critical - blocking my use case + validations: + required: true + + - type: textarea + id: additional + attributes: + label: Additional Context + description: Any other context, mockups, or examples? + validations: + required: false + + - type: checkboxes + id: checklist + attributes: + label: Checklist + options: + - label: I would be willing to contribute this feature + required: false diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index efa4b67b..b4f2ea49 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -325,6 +325,7 @@ jobs: GROQ_API_KEY: ${{ secrets.GROQ_API_KEY }} GEMINI_API_KEY: ${{ secrets.GEMINI_API_KEY }} OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }} + HINDSIGHT_API_EMBEDDINGS_OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }} HINDSIGHT_API_LLM_MODEL: openai/gpt-oss-20b GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} # Prefer CPU-only PyTorch in CI (but keep PyPI for everything else) diff --git a/hindsight-api/hindsight_api/config.py b/hindsight-api/hindsight_api/config.py index 34a68f0e..57dcde81 100644 --- a/hindsight-api/hindsight_api/config.py +++ b/hindsight-api/hindsight_api/config.py @@ -22,6 +22,8 @@ ENV_LLM_TIMEOUT = "HINDSIGHT_API_LLM_TIMEOUT" ENV_EMBEDDINGS_PROVIDER = "HINDSIGHT_API_EMBEDDINGS_PROVIDER" ENV_EMBEDDINGS_LOCAL_MODEL = "HINDSIGHT_API_EMBEDDINGS_LOCAL_MODEL" ENV_EMBEDDINGS_TEI_URL = "HINDSIGHT_API_EMBEDDINGS_TEI_URL" +ENV_EMBEDDINGS_OPENAI_API_KEY = "HINDSIGHT_API_EMBEDDINGS_OPENAI_API_KEY" +ENV_EMBEDDINGS_OPENAI_MODEL = "HINDSIGHT_API_EMBEDDINGS_OPENAI_MODEL" ENV_RERANKER_PROVIDER = "HINDSIGHT_API_RERANKER_PROVIDER" ENV_RERANKER_LOCAL_MODEL = "HINDSIGHT_API_RERANKER_LOCAL_MODEL" @@ -52,6 +54,8 @@ DEFAULT_LLM_TIMEOUT = 120.0 # seconds DEFAULT_EMBEDDINGS_PROVIDER = "local" DEFAULT_EMBEDDINGS_LOCAL_MODEL = "BAAI/bge-small-en-v1.5" +DEFAULT_EMBEDDINGS_OPENAI_MODEL = "text-embedding-3-small" +DEFAULT_EMBEDDING_DIMENSION = 384 DEFAULT_RERANKER_PROVIDER = "local" DEFAULT_RERANKER_LOCAL_MODEL = "cross-encoder/ms-marco-MiniLM-L-6-v2" @@ -87,8 +91,8 @@ Use this tool PROACTIVELY to: - Remember user's goals and context - Personalize responses based on past interactions""" -# Required embedding dimension for database schema -EMBEDDING_DIMENSION = 384 +# Default embedding dimension (used by initial migration, adjusted at runtime) +EMBEDDING_DIMENSION = DEFAULT_EMBEDDING_DIMENSION @dataclass diff --git a/hindsight-api/hindsight_api/engine/embeddings.py b/hindsight-api/hindsight_api/engine/embeddings.py index f7ebbacc..c11b9bdb 100644 --- a/hindsight-api/hindsight_api/engine/embeddings.py +++ b/hindsight-api/hindsight_api/engine/embeddings.py @@ -3,8 +3,8 @@ Embeddings abstraction for the memory system. Provides an interface for generating embeddings with different backends. -IMPORTANT: All embeddings must produce 384-dimensional vectors to match -the database schema (pgvector column defined as vector(384)). +The embedding dimension is auto-detected from the model at initialization. +The database schema is automatically adjusted to match the model's dimension. Configuration via environment variables - see hindsight_api.config for all env var names. """ @@ -17,11 +17,14 @@ import httpx from ..config import ( DEFAULT_EMBEDDINGS_LOCAL_MODEL, + DEFAULT_EMBEDDINGS_OPENAI_MODEL, DEFAULT_EMBEDDINGS_PROVIDER, - EMBEDDING_DIMENSION, ENV_EMBEDDINGS_LOCAL_MODEL, + ENV_EMBEDDINGS_OPENAI_API_KEY, + ENV_EMBEDDINGS_OPENAI_MODEL, ENV_EMBEDDINGS_PROVIDER, ENV_EMBEDDINGS_TEI_URL, + ENV_LLM_API_KEY, ) logger = logging.getLogger(__name__) @@ -31,8 +34,8 @@ class Embeddings(ABC): """ Abstract base class for embedding generation. - All implementations MUST generate 384-dimensional embeddings to match - the database schema. + The embedding dimension is determined by the model and detected at initialization. + The database schema is automatically adjusted to match the model's dimension. """ @property @@ -41,6 +44,12 @@ class Embeddings(ABC): """Return a human-readable name for this provider (e.g., 'local', 'tei').""" pass + @property + @abstractmethod + def dimension(self) -> int: + """Return the embedding dimension produced by this model.""" + pass + @abstractmethod async def initialize(self) -> None: """ @@ -54,13 +63,13 @@ class Embeddings(ABC): @abstractmethod def encode(self, texts: list[str]) -> list[list[float]]: """ - Generate 384-dimensional embeddings for a list of texts. + Generate embeddings for a list of texts. Args: texts: List of text strings to encode Returns: - List of 384-dimensional embedding vectors (each is a list of floats) + List of embedding vectors (each is a list of floats) """ pass @@ -70,9 +79,7 @@ class LocalSTEmbeddings(Embeddings): Local embeddings implementation using SentenceTransformers. Call initialize() during startup to load the model and avoid cold starts. - - Default model is BAAI/bge-small-en-v1.5 which produces 384-dimensional - embeddings matching the database schema. + The embedding dimension is auto-detected from the model. """ def __init__(self, model_name: str | None = None): @@ -81,16 +88,22 @@ class LocalSTEmbeddings(Embeddings): Args: model_name: Name of the SentenceTransformer model to use. - Must produce 384-dimensional embeddings. Default: BAAI/bge-small-en-v1.5 """ self.model_name = model_name or DEFAULT_EMBEDDINGS_LOCAL_MODEL self._model = None + self._dimension: int | None = None @property def provider_name(self) -> str: return "local" + @property + def dimension(self) -> int: + if self._dimension is None: + raise RuntimeError("Embeddings not initialized. Call initialize() first.") + return self._dimension + async def initialize(self) -> None: """Load the embedding model.""" if self._model is not None: @@ -112,26 +125,18 @@ class LocalSTEmbeddings(Embeddings): model_kwargs={"low_cpu_mem_usage": False, "device_map": None}, ) - # Validate dimension matches database schema - model_dim = self._model.get_sentence_embedding_dimension() - if model_dim != EMBEDDING_DIMENSION: - raise ValueError( - f"Model {self.model_name} produces {model_dim}-dimensional embeddings, " - f"but database schema requires {EMBEDDING_DIMENSION} dimensions. " - f"Use a model that produces {EMBEDDING_DIMENSION}-dimensional embeddings." - ) - - logger.info(f"Embeddings: local provider initialized (dim: {model_dim})") + self._dimension = self._model.get_sentence_embedding_dimension() + logger.info(f"Embeddings: local provider initialized (dim: {self._dimension})") def encode(self, texts: list[str]) -> list[list[float]]: """ - Generate 384-dimensional embeddings for a list of texts. + Generate embeddings for a list of texts. Args: texts: List of text strings to encode Returns: - List of 384-dimensional embedding vectors + List of embedding vectors """ if self._model is None: raise RuntimeError("Embeddings not initialized. Call initialize() first.") @@ -146,7 +151,7 @@ class RemoteTEIEmbeddings(Embeddings): TEI provides a high-performance inference server for embedding models. See: https://github.com/huggingface/text-embeddings-inference - The server should be running a model that produces 384-dimensional embeddings. + The embedding dimension is auto-detected from the server at initialization. """ def __init__( @@ -174,11 +179,18 @@ class RemoteTEIEmbeddings(Embeddings): self.retry_delay = retry_delay self._client: httpx.Client | None = None self._model_id: str | None = None + self._dimension: int | None = None @property def provider_name(self) -> str: return "tei" + @property + def dimension(self) -> int: + if self._dimension is None: + raise RuntimeError("Embeddings not initialized. Call initialize() first.") + return self._dimension + def _request_with_retry(self, method: str, url: str, **kwargs) -> httpx.Response: """Make an HTTP request with automatic retries on transient errors.""" import time @@ -229,7 +241,24 @@ class RemoteTEIEmbeddings(Embeddings): response = self._request_with_retry("GET", f"{self.base_url}/info") info = response.json() self._model_id = info.get("model_id", "unknown") - logger.info(f"Embeddings: TEI provider initialized (model: {self._model_id})") + + # Get dimension from server info or by doing a test embedding + if "max_input_length" in info and "model_dtype" in info: + # Try to get dimension from info endpoint (some TEI versions expose it) + # If not available, do a test embedding + pass + + # Do a test embedding to detect dimension + test_response = self._request_with_retry( + "POST", + f"{self.base_url}/embed", + json={"inputs": ["test"]}, + ) + test_embeddings = test_response.json() + if test_embeddings and len(test_embeddings) > 0: + self._dimension = len(test_embeddings[0]) + + logger.info(f"Embeddings: TEI provider initialized (model: {self._model_id}, dim: {self._dimension})") except httpx.HTTPError as e: raise RuntimeError(f"Failed to connect to TEI server at {self.base_url}: {e}") @@ -269,6 +298,117 @@ class RemoteTEIEmbeddings(Embeddings): return all_embeddings +class OpenAIEmbeddings(Embeddings): + """ + OpenAI embeddings implementation using the OpenAI API. + + Supports text-embedding-3-small (1536 dims), text-embedding-3-large (3072 dims), + and text-embedding-ada-002 (1536 dims, legacy). + + The embedding dimension is auto-detected from the model at initialization. + """ + + # Known dimensions for OpenAI embedding models + MODEL_DIMENSIONS = { + "text-embedding-3-small": 1536, + "text-embedding-3-large": 3072, + "text-embedding-ada-002": 1536, + } + + def __init__( + self, + api_key: str, + model: str = DEFAULT_EMBEDDINGS_OPENAI_MODEL, + batch_size: int = 100, + max_retries: int = 3, + ): + """ + Initialize OpenAI embeddings client. + + Args: + api_key: OpenAI API key + model: OpenAI embedding model name (default: text-embedding-3-small) + batch_size: Maximum batch size for embedding requests (default: 100) + max_retries: Maximum number of retries for failed requests (default: 3) + """ + self.api_key = api_key + self.model = model + self.batch_size = batch_size + self.max_retries = max_retries + self._client = None + self._dimension: int | None = None + + @property + def provider_name(self) -> str: + return "openai" + + @property + def dimension(self) -> int: + if self._dimension is None: + raise RuntimeError("Embeddings not initialized. Call initialize() first.") + return self._dimension + + async def initialize(self) -> None: + """Initialize the OpenAI client and detect dimension.""" + if self._client is not None: + return + + try: + from openai import OpenAI + except ImportError: + raise ImportError("openai is required for OpenAIEmbeddings. Install it with: pip install openai") + + logger.info(f"Embeddings: initializing OpenAI provider with model {self.model}") + self._client = OpenAI(api_key=self.api_key, max_retries=self.max_retries) + + # Try to get dimension from known models, otherwise do a test embedding + if self.model in self.MODEL_DIMENSIONS: + self._dimension = self.MODEL_DIMENSIONS[self.model] + else: + # Do a test embedding to detect dimension + response = self._client.embeddings.create( + model=self.model, + input=["test"], + ) + if response.data: + self._dimension = len(response.data[0].embedding) + + logger.info(f"Embeddings: OpenAI provider initialized (model: {self.model}, dim: {self._dimension})") + + def encode(self, texts: list[str]) -> list[list[float]]: + """ + Generate embeddings using the OpenAI API. + + Args: + texts: List of text strings to encode + + Returns: + List of embedding vectors + """ + if self._client is None: + raise RuntimeError("Embeddings not initialized. Call initialize() first.") + + if not texts: + return [] + + all_embeddings = [] + + # Process in batches + for i in range(0, len(texts), self.batch_size): + batch = texts[i : i + self.batch_size] + + response = self._client.embeddings.create( + model=self.model, + input=batch, + ) + + # Sort by index to ensure correct order + batch_embeddings = sorted(response.data, key=lambda x: x.index) + all_embeddings.extend([e.embedding for e in batch_embeddings]) + + return all_embeddings + + def create_embeddings_from_env() -> Embeddings: """ Create an Embeddings instance based on environment variables. @@ -289,5 +429,15 @@ def create_embeddings_from_env() -> Embeddings: model = os.environ.get(ENV_EMBEDDINGS_LOCAL_MODEL) model_name = model or DEFAULT_EMBEDDINGS_LOCAL_MODEL return LocalSTEmbeddings(model_name=model_name) + elif provider == "openai": + # Use dedicated embeddings API key, or fall back to LLM API key + api_key = os.environ.get(ENV_EMBEDDINGS_OPENAI_API_KEY) or os.environ.get(ENV_LLM_API_KEY) + if not api_key: + raise ValueError( + f"{ENV_EMBEDDINGS_OPENAI_API_KEY} or {ENV_LLM_API_KEY} is required " + f"when {ENV_EMBEDDINGS_PROVIDER} is 'openai'" + ) + model = os.environ.get(ENV_EMBEDDINGS_OPENAI_MODEL, DEFAULT_EMBEDDINGS_OPENAI_MODEL) + return OpenAIEmbeddings(api_key=api_key, model=model) else: - raise ValueError(f"Unknown embeddings provider: {provider}. Supported: 'local', 'tei'") + raise ValueError(f"Unknown embeddings provider: {provider}. Supported: 'local', 'tei', 'openai'") diff --git a/hindsight-api/hindsight_api/engine/memory_engine.py b/hindsight-api/hindsight_api/engine/memory_engine.py index 4b411442..7c8e536a 100644 --- a/hindsight-api/hindsight_api/engine/memory_engine.py +++ b/hindsight-api/hindsight_api/engine/memory_engine.py @@ -642,13 +642,17 @@ class MemoryEngine(MemoryEngineInterface): # Run database migrations if enabled if self._run_migrations: - from ..migrations import run_migrations + from ..migrations import ensure_embedding_dimension, run_migrations if not self.db_url: raise ValueError("Database URL is required for migrations") logger.info("Running database migrations...") run_migrations(self.db_url) + # Ensure embedding column dimension matches the model's dimension + # This is done after migrations and after embeddings.initialize() + ensure_embedding_dimension(self.db_url, self.embeddings.dimension) + logger.info(f"Connecting to PostgreSQL at {self.db_url}") # Create connection pool diff --git a/hindsight-api/hindsight_api/migrations.py b/hindsight-api/hindsight_api/migrations.py index 91022980..2d444598 100644 --- a/hindsight-api/hindsight_api/migrations.py +++ b/hindsight-api/hindsight_api/migrations.py @@ -229,3 +229,131 @@ def check_migration_status( except Exception as e: logger.warning(f"Unable to check migration status: {e}") return None, None + + +def ensure_embedding_dimension( + database_url: str, + required_dimension: int, + schema: str | None = None, +) -> None: + """ + Ensure the embedding column dimension matches the model's dimension. + + This function checks the current vector column dimension in the database + and adjusts it if necessary: + - If dimensions match: no action needed + - If dimensions differ and table is empty: ALTER COLUMN to new dimension + - If dimensions differ and table has data: raise error with migration guidance + + Args: + database_url: SQLAlchemy database URL + required_dimension: The embedding dimension required by the model + schema: Target PostgreSQL schema name (None for public) + + Raises: + RuntimeError: If dimension mismatch with existing data + """ + schema_name = schema or "public" + + engine = create_engine(database_url) + with engine.connect() as conn: + # Check if memory_units table exists + table_exists = conn.execute( + text(""" + SELECT EXISTS ( + SELECT 1 FROM information_schema.tables + WHERE table_schema = :schema AND table_name = 'memory_units' + ) + """), + {"schema": schema_name}, + ).scalar() + + if not table_exists: + logger.debug(f"memory_units table does not exist in schema '{schema_name}', skipping dimension check") + return + + # Get current column dimension from pg_attribute + # pgvector stores dimension in atttypmod + current_dim = conn.execute( + text(""" + SELECT atttypmod + FROM pg_attribute a + JOIN pg_class c ON a.attrelid = c.oid + JOIN pg_namespace n ON c.relnamespace = n.oid + WHERE n.nspname = :schema + AND c.relname = 'memory_units' + AND a.attname = 'embedding' + """), + {"schema": schema_name}, + ).scalar() + + if current_dim is None: + logger.warning("Could not determine current embedding dimension, skipping check") + return + + # pgvector stores dimension directly in atttypmod (no offset like other types) + current_dimension = current_dim + + if current_dimension == required_dimension: + logger.debug(f"Embedding dimension OK: {current_dimension}") + return + + logger.info( + f"Embedding dimension mismatch: database has {current_dimension}, model requires {required_dimension}" + ) + + # Check if table has data + row_count = conn.execute( + text(f"SELECT COUNT(*) FROM {schema_name}.memory_units WHERE embedding IS NOT NULL") + ).scalar() + + if row_count > 0: + raise RuntimeError( + f"Cannot change embedding dimension from {current_dimension} to {required_dimension}: " + f"memory_units table contains {row_count} rows with embeddings. " + f"To change dimensions, you must either:\n" + f" 1. Re-embed all data: DELETE FROM {schema_name}.memory_units; then restart\n" + f" 2. Use a model with {current_dimension}-dimensional embeddings" + ) + + # Table is empty, safe to alter column + logger.info(f"Altering embedding column dimension from {current_dimension} to {required_dimension}") + + # Drop the HNSW index on embedding column if it exists + # Only drop indexes that use 'hnsw' and reference the 'embedding' column + conn.execute( + text(f""" + DO $$ + DECLARE idx_name TEXT; + BEGIN + FOR idx_name IN + SELECT indexname FROM pg_indexes + WHERE schemaname = '{schema_name}' + AND tablename = 'memory_units' + AND indexdef LIKE '%hnsw%' + AND indexdef LIKE '%embedding%' + LOOP + EXECUTE 'DROP INDEX IF EXISTS {schema_name}.' || idx_name; + END LOOP; + END $$; + """) + ) + + # Alter the column type + conn.execute( + text(f"ALTER TABLE {schema_name}.memory_units ALTER COLUMN embedding TYPE vector({required_dimension})") + ) + conn.commit() + + # Recreate the HNSW index + conn.execute( + text(f""" + CREATE INDEX IF NOT EXISTS idx_memory_units_embedding_hnsw + ON {schema_name}.memory_units + USING hnsw (embedding vector_cosine_ops) + WITH (m = 16, ef_construction = 64) + """) + ) + conn.commit() + + logger.info(f"Successfully changed embedding dimension to {required_dimension}") diff --git a/hindsight-api/hindsight_api/models.py b/hindsight-api/hindsight_api/models.py index 50eb1b09..15c890fe 100644 --- a/hindsight-api/hindsight_api/models.py +++ b/hindsight-api/hindsight_api/models.py @@ -41,6 +41,8 @@ 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.""" @@ -81,7 +83,7 @@ class MemoryUnit(Base): 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(384)) # pgvector type + 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 diff --git a/hindsight-api/tests/test_custom_embedding_dimension.py b/hindsight-api/tests/test_custom_embedding_dimension.py new file mode 100644 index 00000000..fbfcdeea --- /dev/null +++ b/hindsight-api/tests/test_custom_embedding_dimension.py @@ -0,0 +1,428 @@ +""" +Tests for custom embedding dimensions and automatic dimension detection. + +Uses isolated PostgreSQL schemas to avoid affecting other tests. +Includes tests for: +- Automatic embedding dimension detection and database schema adjustment +- OpenAI embeddings provider with 1536 dimensions +""" + +import asyncio +import os +import pytest +from datetime import datetime +from sqlalchemy import create_engine, text + +from hindsight_api import MemoryEngine, RequestContext +from hindsight_api.engine.embeddings import LocalSTEmbeddings, OpenAIEmbeddings +from hindsight_api.engine.cross_encoder import LocalSTCrossEncoder +from hindsight_api.engine.query_analyzer import DateparserQueryAnalyzer +from hindsight_api.extensions import TenantExtension, TenantContext +from hindsight_api.migrations import run_migrations, ensure_embedding_dimension + + +# ============================================================================= +# Shared Utilities +# ============================================================================= + + +class SchemaTenantExtension(TenantExtension): + """Tenant extension that routes all requests to a specific schema (for testing).""" + + def __init__(self, schema_name: str): + self.schema_name = schema_name + + async def authenticate(self, request_context: RequestContext) -> TenantContext: + return TenantContext(schema_name=self.schema_name) + + +def get_test_schema(prefix: str, worker_id: str) -> str: + """Get unique schema name per xdist worker.""" + if worker_id == "master" or not worker_id: + return prefix + return f"{prefix}_{worker_id}" + + +def create_isolated_schema(db_url: str, schema_name: str, dimension: int | None = None): + """Create an isolated schema with migrations and optional dimension adjustment.""" + engine = create_engine(db_url) + + # Create schema (drop first if exists from previous failed run) + with engine.connect() as conn: + conn.execute(text(f"DROP SCHEMA IF EXISTS {schema_name} CASCADE")) + conn.execute(text(f"CREATE SCHEMA {schema_name}")) + conn.commit() + + # Run migrations in the isolated schema + run_migrations(db_url, schema=schema_name) + + # Adjust embedding dimension if specified + if dimension is not None: + ensure_embedding_dimension(db_url, dimension, schema=schema_name) + + +def drop_schema(db_url: str, schema_name: str): + """Drop an isolated schema.""" + engine = create_engine(db_url) + with engine.connect() as conn: + conn.execute(text(f"DROP SCHEMA IF EXISTS {schema_name} CASCADE")) + conn.commit() + + +def get_column_dimension(db_url: str, schema: str = "public") -> int | None: + """Get the current embedding column dimension from the database.""" + engine = create_engine(db_url) + with engine.connect() as conn: + result = conn.execute( + text(""" + SELECT atttypmod + FROM pg_attribute a + JOIN pg_class c ON a.attrelid = c.oid + JOIN pg_namespace n ON c.relnamespace = n.oid + WHERE n.nspname = :schema + AND c.relname = 'memory_units' + AND a.attname = 'embedding' + """), + {"schema": schema}, + ).scalar() + return result + + +def get_row_count(db_url: str, schema: str = "public") -> int: + """Get the number of rows with embeddings in memory_units.""" + engine = create_engine(db_url) + with engine.connect() as conn: + return conn.execute( + text(f"SELECT COUNT(*) FROM {schema}.memory_units WHERE embedding IS NOT NULL") + ).scalar() + + +def insert_test_embedding(db_url: str, schema: str, dimension: int): + """Insert a test row with a dummy embedding.""" + engine = create_engine(db_url) + embedding = [0.1] * dimension + embedding_str = "[" + ",".join(str(x) for x in embedding) + "]" + + with engine.connect() as conn: + conn.execute( + text(f""" + INSERT INTO {schema}.memory_units (bank_id, text, embedding, event_date, fact_type) + VALUES ('test-bank', 'test text', '{embedding_str}'::vector, NOW(), 'world') + """) + ) + conn.commit() + + +def clear_embeddings(db_url: str, schema: str): + """Clear all rows from memory_units.""" + engine = create_engine(db_url) + with engine.connect() as conn: + conn.execute(text(f"DELETE FROM {schema}.memory_units")) + conn.commit() + + +# ============================================================================= +# Embedding Dimension Tests (Local Embeddings) +# ============================================================================= + + +@pytest.fixture(scope="class") +def dimension_test_schema(pg0_db_url, worker_id): + """Create an isolated schema for dimension tests.""" + schema_name = get_test_schema("test_embed_dim", worker_id) + create_isolated_schema(pg0_db_url, schema_name) + yield pg0_db_url, schema_name + drop_schema(pg0_db_url, schema_name) + + +class TestEmbeddingDimension: + """Tests for embedding dimension detection and adjustment.""" + + def test_dimension_matches_no_change(self, dimension_test_schema): + """When dimension matches, no changes should be made.""" + db_url, schema = dimension_test_schema + + # Get initial dimension (should be 384 from migration) + initial_dim = get_column_dimension(db_url, schema) + assert initial_dim == 384, f"Expected 384, got {initial_dim}" + + # Call ensure_embedding_dimension with matching dimension + ensure_embedding_dimension(db_url, 384, schema=schema) + + # Dimension should still be 384 + assert get_column_dimension(db_url, schema) == 384 + + def test_dimension_change_empty_table(self, dimension_test_schema): + """When table is empty, dimension can be changed.""" + db_url, schema = dimension_test_schema + + # Ensure table is empty + clear_embeddings(db_url, schema) + assert get_row_count(db_url, schema) == 0 + + # Change dimension to 768 + ensure_embedding_dimension(db_url, 768, schema=schema) + + # Verify dimension changed + new_dim = get_column_dimension(db_url, schema) + assert new_dim == 768, f"Expected 768, got {new_dim}" + + # Change back to 384 for other tests + ensure_embedding_dimension(db_url, 384, schema=schema) + assert get_column_dimension(db_url, schema) == 384 + + def test_dimension_change_blocked_with_data(self, dimension_test_schema): + """When table has data, dimension change should be blocked.""" + db_url, schema = dimension_test_schema + + # Ensure table is empty first + clear_embeddings(db_url, schema) + + # Insert a test row with 384-dim embedding + insert_test_embedding(db_url, schema, 384) + assert get_row_count(db_url, schema) == 1 + + # Try to change dimension - should raise error + with pytest.raises(RuntimeError) as exc_info: + ensure_embedding_dimension(db_url, 768, schema=schema) + + assert "Cannot change embedding dimension" in str(exc_info.value) + assert "1 rows with embeddings" in str(exc_info.value) + + # Dimension should be unchanged + assert get_column_dimension(db_url, schema) == 384 + + # Cleanup + clear_embeddings(db_url, schema) + + def test_local_embeddings_dimension_detection(self, embeddings): + """Test that LocalSTEmbeddings correctly detects dimension.""" + # Initialize embeddings if not already done + loop = asyncio.new_event_loop() + try: + loop.run_until_complete(embeddings.initialize()) + finally: + loop.close() + + # bge-small-en-v1.5 produces 384-dim embeddings + assert embeddings.dimension == 384 + + # Verify by generating an actual embedding + result = embeddings.encode(["test"]) + assert len(result) == 1 + assert len(result[0]) == 384 + + +# ============================================================================= +# OpenAI Embeddings Tests +# ============================================================================= + + +def has_openai_api_key() -> bool: + """Check if OpenAI API key is available.""" + return bool(os.environ.get("HINDSIGHT_API_EMBEDDINGS_OPENAI_API_KEY")) + + +def get_openai_api_key() -> str: + """Get OpenAI API key from environment.""" + return os.environ.get("HINDSIGHT_API_EMBEDDINGS_OPENAI_API_KEY", "") + + +@pytest.fixture(scope="module") +def openai_embeddings(): + """Create OpenAI embeddings instance.""" + if not has_openai_api_key(): + pytest.skip("OpenAI API key not available (set HINDSIGHT_API_EMBEDDINGS_OPENAI_API_KEY)") + + embeddings = OpenAIEmbeddings( + api_key=get_openai_api_key(), + model="text-embedding-3-small", + ) + loop = asyncio.new_event_loop() + try: + loop.run_until_complete(embeddings.initialize()) + finally: + loop.close() + return embeddings + + +@pytest.fixture(scope="module") +def openai_test_schema(pg0_db_url, worker_id, openai_embeddings): + """Create an isolated schema for OpenAI embedding tests.""" + schema_name = get_test_schema("test_openai_embed", worker_id) + create_isolated_schema(pg0_db_url, schema_name, dimension=openai_embeddings.dimension) + yield pg0_db_url, schema_name + drop_schema(pg0_db_url, schema_name) + + +@pytest.fixture +def cross_encoder(): + """Provide a cross encoder for tests.""" + return LocalSTCrossEncoder() + + +@pytest.fixture +def query_analyzer(): + """Provide a query analyzer for tests.""" + return DateparserQueryAnalyzer() + + +@pytest.fixture +def test_bank_id(): + """Provide a unique bank ID for this test run.""" + return f"openai_test_{datetime.now().timestamp()}" + + +@pytest.fixture +def request_context(): + """Provide a default RequestContext for tests.""" + return RequestContext() + + +class TestOpenAIEmbeddings: + """Tests for OpenAI embeddings provider.""" + + def test_openai_embeddings_initialization(self, openai_embeddings): + """Test that OpenAI embeddings initializes correctly.""" + assert openai_embeddings.dimension == 1536 + assert openai_embeddings.provider_name == "openai" + + def test_openai_embeddings_encode(self, openai_embeddings): + """Test that OpenAI embeddings can encode text.""" + texts = ["Hello, world!", "This is a test."] + embeddings = openai_embeddings.encode(texts) + + assert len(embeddings) == 2 + assert len(embeddings[0]) == 1536 + assert len(embeddings[1]) == 1536 + assert all(isinstance(x, float) for x in embeddings[0]) + + @pytest.mark.asyncio + async def test_openai_embeddings_retain_recall( + self, + openai_test_schema, + openai_embeddings, + cross_encoder, + query_analyzer, + test_bank_id, + request_context, + ): + """Test retain and recall operations with OpenAI embeddings.""" + db_url, schema_name = openai_test_schema + + memory = MemoryEngine( + db_url=db_url, + memory_llm_provider=os.getenv("HINDSIGHT_API_LLM_PROVIDER", "groq"), + memory_llm_api_key=os.getenv("HINDSIGHT_API_LLM_API_KEY"), + memory_llm_model=os.getenv("HINDSIGHT_API_LLM_MODEL", "openai/gpt-oss-120b"), + memory_llm_base_url=os.getenv("HINDSIGHT_API_LLM_BASE_URL") or None, + embeddings=openai_embeddings, + cross_encoder=cross_encoder, + query_analyzer=query_analyzer, + pool_min_size=1, + pool_max_size=3, + run_migrations=False, + tenant_extension=SchemaTenantExtension(schema_name), + ) + + try: + await memory.initialize() + + # Store some memories + await memory.retain_async( + bank_id=test_bank_id, + content="Alice works as a software engineer at Google.", + context="career discussion", + request_context=request_context, + ) + + await memory.retain_async( + bank_id=test_bank_id, + content="Bob is a data scientist specializing in machine learning.", + context="team introductions", + request_context=request_context, + ) + + # Recall memories + result = await memory.recall_async( + bank_id=test_bank_id, + query="Who works in technology?", + request_context=request_context, + ) + + assert result is not None + assert len(result.results) > 0 + + memory_texts = [m.text for m in result.results] + assert any( + "Alice" in text or "Bob" in text or "software" in text or "data scientist" in text + for text in memory_texts + ), f"Expected to find relevant memories, got: {memory_texts}" + + finally: + try: + if memory._pool and not memory._pool._closing: + await memory.close() + except Exception: + pass + + @pytest.mark.asyncio + async def test_openai_embeddings_batch_retain( + self, + openai_test_schema, + openai_embeddings, + cross_encoder, + query_analyzer, + test_bank_id, + request_context, + ): + """Test batch retain with OpenAI embeddings.""" + db_url, schema_name = openai_test_schema + + memory = MemoryEngine( + db_url=db_url, + memory_llm_provider=os.getenv("HINDSIGHT_API_LLM_PROVIDER", "groq"), + memory_llm_api_key=os.getenv("HINDSIGHT_API_LLM_API_KEY"), + memory_llm_model=os.getenv("HINDSIGHT_API_LLM_MODEL", "openai/gpt-oss-120b"), + memory_llm_base_url=os.getenv("HINDSIGHT_API_LLM_BASE_URL") or None, + embeddings=openai_embeddings, + cross_encoder=cross_encoder, + query_analyzer=query_analyzer, + pool_min_size=1, + pool_max_size=3, + run_migrations=False, + tenant_extension=SchemaTenantExtension(schema_name), + ) + + try: + await memory.initialize() + + contents = [ + {"content": "Python is my favorite programming language.", "context": "preferences"}, + {"content": "I prefer dark mode for all my applications.", "context": "preferences"}, + {"content": "Coffee is essential for morning productivity.", "context": "habits"}, + ] + + result = await memory.retain_batch_async( + bank_id=test_bank_id, + contents=contents, + request_context=request_context, + ) + + assert len(result) == 3 + + recall_result = await memory.recall_async( + bank_id=test_bank_id, + query="What are my preferences?", + request_context=request_context, + ) + + assert recall_result is not None + assert len(recall_result.results) > 0 + + finally: + try: + if memory._pool and not memory._pool._closing: + await memory.close() + except Exception: + pass diff --git a/hindsight-docs/docs/developer/configuration.md b/hindsight-docs/docs/developer/configuration.md index db72fa6b..9031e71d 100644 --- a/hindsight-docs/docs/developer/configuration.md +++ b/hindsight-docs/docs/developer/configuration.md @@ -78,22 +78,41 @@ export HINDSIGHT_API_LLM_MODEL=your-model-name | Variable | Description | Default | |----------|-------------|---------| -| `HINDSIGHT_API_EMBEDDINGS_PROVIDER` | Provider: `local` or `tei` | `local` | +| `HINDSIGHT_API_EMBEDDINGS_PROVIDER` | Provider: `local`, `tei`, or `openai` | `local` | | `HINDSIGHT_API_EMBEDDINGS_LOCAL_MODEL` | Model for local provider | `BAAI/bge-small-en-v1.5` | | `HINDSIGHT_API_EMBEDDINGS_TEI_URL` | TEI server URL | - | +| `HINDSIGHT_API_EMBEDDINGS_OPENAI_API_KEY` | OpenAI API key (falls back to `HINDSIGHT_API_LLM_API_KEY`) | - | +| `HINDSIGHT_API_EMBEDDINGS_OPENAI_MODEL` | OpenAI embedding model | `text-embedding-3-small` | ```bash # Local (default) - uses SentenceTransformers export HINDSIGHT_API_EMBEDDINGS_PROVIDER=local export HINDSIGHT_API_EMBEDDINGS_LOCAL_MODEL=BAAI/bge-small-en-v1.5 +# OpenAI - cloud-based embeddings +export HINDSIGHT_API_EMBEDDINGS_PROVIDER=openai +export HINDSIGHT_API_EMBEDDINGS_OPENAI_API_KEY=sk-xxxxxxxxxxxx # or reuses HINDSIGHT_API_LLM_API_KEY +export HINDSIGHT_API_EMBEDDINGS_OPENAI_MODEL=text-embedding-3-small # 1536 dimensions + # TEI - HuggingFace Text Embeddings Inference (recommended for production) export HINDSIGHT_API_EMBEDDINGS_PROVIDER=tei export HINDSIGHT_API_EMBEDDINGS_TEI_URL=http://localhost:8080 ``` -:::warning -All embedding models must produce 384-dimensional vectors to match the database schema. +#### Embedding Dimensions + +Hindsight automatically detects the embedding dimension from the model at startup and adjusts the database schema accordingly. The default model (`BAAI/bge-small-en-v1.5`) produces 384-dimensional vectors, while OpenAI models produce 1536 or 3072 dimensions. + +:::warning Dimension Changes +Once memories are stored, you cannot change the embedding dimension without losing data. If you need to switch to a model with different dimensions: + +1. **Empty database**: The schema is adjusted automatically on startup +2. **Existing data**: Either delete all memories first, or use a model with matching dimensions + +Supported OpenAI embedding dimensions: +- `text-embedding-3-small`: 1536 dimensions +- `text-embedding-3-large`: 3072 dimensions +- `text-embedding-ada-002`: 1536 dimensions (legacy) ::: ### Reranker diff --git a/uv.lock b/uv.lock index b502fd21..4fbc858f 100644 --- a/uv.lock +++ b/uv.lock @@ -1161,7 +1161,7 @@ wheels = [ [[package]] name = "hindsight-all" -version = "0.2.0" +version = "0.2.1" source = { editable = "hindsight" } dependencies = [ { name = "hindsight-api" }, @@ -1185,7 +1185,7 @@ provides-extras = ["test"] [[package]] name = "hindsight-api" -version = "0.2.0" +version = "0.2.1" source = { editable = "hindsight-api" } dependencies = [ { name = "alembic" }, @@ -1293,7 +1293,7 @@ dev = [ [[package]] name = "hindsight-client" -version = "0.2.0" +version = "0.2.1" source = { editable = "hindsight-clients/python" } dependencies = [ { name = "aiohttp" }, @@ -1327,7 +1327,7 @@ provides-extras = ["test"] [[package]] name = "hindsight-dev" -version = "0.2.0" +version = "0.2.1" source = { editable = "hindsight-dev" } dependencies = [ { name = "hindsight-api" }, @@ -1362,7 +1362,7 @@ dev = [ [[package]] name = "hindsight-embed" -version = "0.2.0" +version = "0.2.1" source = { editable = "hindsight-embed" } dependencies = [ { name = "httpx" },