Compare commits

...
Author SHA1 Message Date
Nicolò Boschi 3347ea6e09 fix tests 2026-01-05 14:34:00 +01:00
Nicolò Boschi 96631b69ef feat: configurable embedding dimensions + OpenAI Embeddings 2026-01-05 14:24:17 +01:00
12 changed files with 935 additions and 38 deletions
+71
View File
@@ -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
+8
View File
@@ -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.
@@ -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
+1
View File
@@ -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)
+6 -2
View File
@@ -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
+176 -26
View File
@@ -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'")
@@ -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
+128
View File
@@ -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}")
+3 -1
View File
@@ -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
@@ -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
+22 -3
View File
@@ -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
Generated
+5 -5
View File
@@ -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" },