Semantic search with pgvector

alembic/versions/20250825_123447_7aa80d34dbdd_add_pgvector.py

@@ -0,0 +1,117 @@
+"""add pgvector
+
+Revision ID: 7aa80d34dbdd
+Revises: 02b804d687ee
+Create Date: 2025-08-25 12:34:47.832367
+
+"""
+
+from typing import Sequence, Union
+import os
+import random
+
+import openai
+from alembic import op
+from pgvector.sqlalchemy import Vector
+import sqlalchemy as sa
+from sqlalchemy import text
+
+# revision identifiers, used by Alembic.
+revision: str = "7aa80d34dbdd"
+down_revision: Union[str, None] = "02b804d687ee"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def generate_embeddings_for_existing_data():
+    """Generate embeddings for existing brain_region, species, and strain data."""
+    # Get connection
+    connection = op.get_bind()
+
+    # Collect all entity data
+    all_entities = []
+
+    # Collect brain region data
+    brain_regions = connection.execute(text("SELECT id, name FROM brain_region")).fetchall()
+    for brain_region in brain_regions:
+        all_entities.append(("brain_region", brain_region.id, brain_region.name))
+
+    # Collect species data
+    species = connection.execute(text("SELECT id, name FROM species")).fetchall()
+    for sp in species:
+        all_entities.append(("species", sp.id, sp.name))
+
+    # Collect strain data
+    strains = connection.execute(text("SELECT id, name FROM strain")).fetchall()
+    for strain in strains:
+        all_entities.append(("strain", strain.id, strain.name))
+
+    # Generate embeddings based on available API key
+    api_key = os.getenv("OPENAI_API_KEY")
+
+    if api_key:
+        # Use OpenAI API for real embeddings
+        client = openai.OpenAI(api_key=api_key)
+
+        # Generate all embeddings in a single API call
+        names = [entity[2] for entity in all_entities]
+        response = client.embeddings.create(model="text-embedding-3-small", input=names)
+
+        # Extract embeddings from response
+        embeddings = [embedding.embedding for embedding in response.data]
+    else:
+        # Use random vectors when OpenAI key is not provided
+        embeddings = []
+        for _ in all_entities:
+            random_embedding = [random.random() for _ in range(1536)]
+            embeddings.append(random_embedding)
+
+    # Update database with generated embeddings (shared logic)
+    for (table_name, entity_id, _), embedding in zip(all_entities, embeddings):
+        # Convert embedding to string format for pgvector
+        embedding_str = str(embedding)
+        connection.execute(
+            text(f"UPDATE {table_name} SET embedding = :embedding WHERE id = :id"),
+            {"embedding": embedding_str, "id": entity_id},
+        )
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    # Enable the pgvector extension
+    op.execute("CREATE EXTENSION IF NOT EXISTS vector;")
+
+    op.add_column(
+        "brain_region",
+        sa.Column("embedding", Vector(dim=1536), nullable=True),
+    )
+
+    op.add_column(
+        "species",
+        sa.Column("embedding", Vector(dim=1536), nullable=True),
+    )
+
+    op.add_column(
+        "strain",
+        sa.Column("embedding", Vector(dim=1536), nullable=True),
+    )
+
+    # Generate embeddings for existing data
+    generate_embeddings_for_existing_data()
+
+    # Now make columns non-nullable
+    op.alter_column("brain_region", "embedding", nullable=False)
+    op.alter_column("species", "embedding", nullable=False)
+    op.alter_column("strain", "embedding", nullable=False)
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_column("strain", "embedding")
+    op.drop_column("species", "embedding")
+    op.drop_column("brain_region", "embedding")
+
+    # Disable the pgvector extension
+    op.execute("DROP EXTENSION IF EXISTS vector;")
+    # ### end Alembic commands ###

app/config.py

@@ -1,6 +1,6 @@
 from typing import Literal
 
-from pydantic import PostgresDsn, field_validator
+from pydantic import PostgresDsn, SecretStr, field_validator
 from pydantic_core.core_schema import ValidationInfo
 from pydantic_settings import BaseSettings, SettingsConfigDict
 
@@ -62,6 +62,8 @@ class Settings(BaseSettings):
     DB_POOL_PRE_PING: bool = False
     DB_MAX_OVERFLOW: int = 0
 
+    OPENAI_API_KEY: SecretStr | None = None
+
     @field_validator("DB_URI", mode="before")
     @classmethod
     def build_db_uri(cls, v: str, info: ValidationInfo) -> str:

app/db/model.py

@@ -4,6 +4,7 @@
 from uuid import UUID
 
 import sqlalchemy as sa
+from pgvector.sqlalchemy import Vector
 from sqlalchemy import (
     BigInteger,
     DateTime,
@@ -175,12 +176,14 @@ class BrainRegion(Identifiable):
     hierarchy_id: Mapped[uuid.UUID] = mapped_column(
         ForeignKey("brain_region_hierarchy.id"), index=True
     )
+    embedding: Mapped[Vector] = mapped_column(Vector(1536), nullable=False)
 
 
 class Species(Identifiable):
     __tablename__ = "species"
     name: Mapped[str] = mapped_column(unique=True, index=True)
     taxonomy_id: Mapped[str] = mapped_column(unique=True, index=True)
+    embedding: Mapped[Vector] = mapped_column(Vector(1536), nullable=False)
 
 
 class Strain(Identifiable):
@@ -189,6 +192,7 @@ class Strain(Identifiable):
     taxonomy_id: Mapped[str] = mapped_column(unique=True, index=True)
     species_id: Mapped[uuid.UUID] = mapped_column(ForeignKey("species.id"), index=True)
     species = relationship("Species", uselist=False)
+    embedding: Mapped[Vector] = mapped_column(Vector(1536), nullable=False)
 
     __table_args__ = (
         # needed for the composite foreign key in SpeciesMixin

app/queries/common.py

@@ -238,6 +238,7 @@ def router_read_many[T: BaseModel, I: Identifiable](  # noqa: PLR0913
     name_to_facet_query_params: dict[str, FacetQueryParams] | None,
     filter_model: CustomFilter[I],
     filter_joins: dict[str, ApplyOperations] | None = None,
+    embedding: list[float] | None = None,
 ) -> ListResponse[T]:
     """Read multiple models from the database.
 
@@ -258,6 +259,7 @@ def router_read_many[T: BaseModel, I: Identifiable](  # noqa: PLR0913
         filter_joins: mapping of filter names to join functions. The keys should match both:
             - the nested filters attributes, to choose which joins should be applied for filtering.
             - the keys in `name_to_facet_query_params`, for retrieving the facets.
+        embedding: optional list of floats representing an embedding vector for semantic search.
 
     Returns:
         the list of model data, pagination, and facets as a Pydantic model.
@@ -291,6 +293,15 @@ def router_read_many[T: BaseModel, I: Identifiable](  # noqa: PLR0913
         .limit(pagination_request.page_size)
     )
 
+    # Add semantic similarity ordering if embedding is provided and model has embedding field
+    if embedding is not None and hasattr(db_model_class, "embedding"):
+        # Remove existing ordering clauses and replace with semantic similarity ordering
+        if getattr(data_query, "_order_by_clauses", None):
+            # Clear existing ordering by setting _order_by_clauses to empty tuple
+            data_query._order_by_clauses = ()  # noqa: SLF001
+
+        data_query = data_query.order_by(db_model_class.embedding.l2_distance(embedding))  # type: ignore[attr-defined]
+
     if apply_data_query_operations:
         data_query = apply_data_query_operations(data_query)
 

app/schemas/species.py

@@ -1,6 +1,7 @@
 from uuid import UUID
 
-from pydantic import BaseModel, ConfigDict
+from pydantic import BaseModel, ConfigDict, Field
+from pydantic.json_schema import SkipJsonSchema
 
 from app.schemas.agent import CreatedByUpdatedByMixin
 from app.schemas.base import CreationMixin, IdentifiableMixin
@@ -10,26 +11,28 @@ class SpeciesCreate(BaseModel):
     model_config = ConfigDict(from_attributes=True)
     name: str
     taxonomy_id: str
+    embedding: SkipJsonSchema[list[float] | None] = None
 
 
 class SpeciesRead(SpeciesCreate, CreationMixin, CreatedByUpdatedByMixin, IdentifiableMixin):
-    pass
+    embedding: SkipJsonSchema[list[float] | None] = Field(default=None, exclude=True)
 
 
 class NestedSpeciesRead(SpeciesCreate, IdentifiableMixin):
-    pass
+    embedding: SkipJsonSchema[list[float] | None] = Field(default=None, exclude=True)
 
 
 class StrainCreate(BaseModel):
     model_config = ConfigDict(from_attributes=True)
     name: str
     taxonomy_id: str
     species_id: UUID
+    embedding: SkipJsonSchema[list[float] | None] = None
 
 
 class StrainRead(StrainCreate, CreationMixin, CreatedByUpdatedByMixin, IdentifiableMixin):
-    pass
+    embedding: SkipJsonSchema[list[float] | None] = Field(default=None, exclude=True)
 
 
 class NestedStrainRead(StrainCreate, IdentifiableMixin):
-    pass
+    embedding: SkipJsonSchema[list[float] | None] = Field(default=None, exclude=True)

app/service/brain_region.py

@@ -10,14 +10,21 @@
 from app.filters.brain_region import BrainRegionFilterDep
 from app.schemas.base import BrainRegionRead
 from app.schemas.types import ListResponse
+from app.utils.embedding import generate_embedding
 
 
 def read_many(
     *,
     db: SessionDep,
     pagination_request: PaginationQuery,
     brain_region_filter: BrainRegionFilterDep,
+    semantic_search: str | None = None,
 ) -> ListResponse[BrainRegionRead]:
+    embedding = None
+
+    if semantic_search is not None:
+        embedding = generate_embedding(semantic_search)
+
     return app.queries.common.router_read_many(
         db=db,
         db_model_class=BrainRegion,
@@ -32,6 +39,7 @@ def read_many(
         response_schema_class=BrainRegionRead,
         name_to_facet_query_params=None,
         filter_model=brain_region_filter,
+        embedding=embedding,
     )
 
 

app/service/species.py

@@ -12,6 +12,7 @@
 from app.queries.factory import query_params_factory
 from app.schemas.species import SpeciesCreate, SpeciesRead
 from app.schemas.types import ListResponse
+from app.utils.embedding import generate_embedding
 
 
 def _load(query: sa.Select):
@@ -43,6 +44,9 @@ def create_one(
     species: SpeciesCreate,
     user_context: AdminContextDep,
 ) -> SpeciesRead:
+    # Generate embedding using OpenAI API
+    species.embedding = generate_embedding(species.name)
+
     return app.queries.common.router_create_one(
         db=db,
         db_model_class=Species,
@@ -58,7 +62,13 @@ def read_many(
     db: SessionDep,
     pagination_request: PaginationQuery,
     species_filter: SpeciesFilterDep,
+    semantic_search: str | None = None,
 ) -> ListResponse[SpeciesRead]:
+    embedding = None
+
+    if semantic_search is not None:
+        embedding = generate_embedding(semantic_search)
+
     facet_keys = filter_keys = [
         "created_by",
         "updated_by",
@@ -84,4 +94,5 @@ def read_many(
         name_to_facet_query_params=name_to_facet_query_params,
         filter_model=species_filter,
         filter_joins=filter_joins,
+        embedding=embedding,
     )

app/service/strain.py

@@ -12,6 +12,7 @@
 from app.queries.factory import query_params_factory
 from app.schemas.species import StrainCreate, StrainRead
 from app.schemas.types import ListResponse
+from app.utils.embedding import generate_embedding
 
 
 def _load(query: sa.Select):
@@ -23,10 +24,17 @@ def _load(query: sa.Select):
 
 
 def read_many(
+    *,
     db: SessionDep,
     pagination_request: PaginationQuery,
     strain_filter: StrainFilterDep,
+    semantic_search: str | None = None,
 ) -> ListResponse[StrainRead]:
+    embedding = None
+
+    if semantic_search is not None:
+        embedding = generate_embedding(semantic_search)
+
     facet_keys = filter_keys = [
         "created_by",
         "updated_by",
@@ -52,6 +60,7 @@ def read_many(
         name_to_facet_query_params=name_to_facet_query_params,
         filter_model=strain_filter,
         filter_joins=filter_joins,
+        embedding=embedding,
     )
 
 
@@ -69,6 +78,9 @@ def read_one(id_: uuid.UUID, db: SessionDep) -> StrainRead:
 def create_one(
     json_model: StrainCreate, db: SessionDep, user_context: AdminContextDep
 ) -> StrainRead:
+    # Generate embedding using OpenAI API
+    json_model.embedding = generate_embedding(json_model.name)
+
     return app.queries.common.router_create_one(
         db=db,
         db_model_class=Strain,

app/utils/embedding.py

@@ -0,0 +1,32 @@
+"""Utility functions for generating embeddings using OpenAI API."""
+
+import openai
+
+from app.config import settings
+
+
+def generate_embedding(text: str, model: str = "text-embedding-3-small") -> list[float]:
+    """Generate an embedding for the given text using OpenAI API.
+
+    Args:
+        text: The text to generate an embedding for
+        model: The OpenAI embedding model to use (default: text-embedding-3-small)
+
+    Returns:
+        A list of floats representing the embedding vector
+
+    Raises:
+        ValueError: If OpenAI API key is not configured
+    """
+    if settings.OPENAI_API_KEY is None:
+        message = "OpenAI API key is not configured."
+        raise ValueError(message)
+
+    openai_api_key = settings.OPENAI_API_KEY.get_secret_value()
+
+    # Generate embedding using OpenAI API
+    client = openai.OpenAI(api_key=openai_api_key)
+    response = client.embeddings.create(model=model, input=text)
+
+    # Return the generated embedding
+    return response.data[0].embedding

docker-compose.run.yml

@@ -29,7 +29,7 @@ services:
 
   db:
     profiles: [run]
-    image: postgres:17-alpine
+    image: pgvector/pgvector:0.8.0-pg17-trixie
     environment:
     - POSTGRES_USER=entitycore
     - POSTGRES_PASSWORD=entitycore