diff --git a/docs/python-integration.md b/docs/python-integration.md index ff67dbd5d..7afc6dd79 100644 --- a/docs/python-integration.md +++ b/docs/python-integration.md @@ -197,7 +197,7 @@ Then you can use the `create_vectorizer` operation to create a vectorizer for yo ```python from alembic import op -from pgai.vectorizer.configuration import ( +from pgai.alembic.configuration import ( OpenAIConfig, CharacterTextSplitterConfig, PythonTemplateConfig diff --git a/projects/extension/Dockerfile b/projects/extension/Dockerfile index 470b297bb..6277e1b8c 100644 --- a/projects/extension/Dockerfile +++ b/projects/extension/Dockerfile @@ -52,6 +52,10 @@ WORKDIR /pgai COPY . . RUN just build install +RUN mkdir -p /docker-entrypoint-initdb.d && \ + echo "#!/bin/bash" > /docker-entrypoint-initdb.d/configure-timescaledb.sh && \ + echo "echo \"shared_preload_libraries = 'timescaledb'\" >> \${PGDATA}/postgresql.conf" >> /docker-entrypoint-initdb.d/configure-timescaledb.sh && \ + chmod +x /docker-entrypoint-initdb.d/configure-timescaledb.sh ############################################################################### # image for use in extension development diff --git a/projects/pgai/pgai/vectorizer/configuration.py b/projects/pgai/pgai/alembic/configuration.py similarity index 98% rename from projects/pgai/pgai/vectorizer/configuration.py rename to projects/pgai/pgai/alembic/configuration.py index 239ac4f4e..d902b272c 100644 --- a/projects/pgai/pgai/vectorizer/configuration.py +++ b/projects/pgai/pgai/alembic/configuration.py @@ -65,7 +65,6 @@ def to_sql_argument(self) -> str: base_name = re.sub("([a-z0-9])([A-Z])", r"\1_\2", base_name).lower() # Remove 'config' and clean up any double underscores base_name = base_name.replace("config", "").strip("_") - # Remove any duplicate underscores that might have been created fn_name = f"{self.arg_type}_{base_name}" return f", {self.arg_type} => ai.{fn_name}({format_sql_params(params)})" # type: ignore diff --git a/projects/pgai/pgai/alembic/operations.py b/projects/pgai/pgai/alembic/operations.py index 686eeed8b..e89ff5070 100644 --- a/projects/pgai/pgai/alembic/operations.py +++ b/projects/pgai/pgai/alembic/operations.py @@ -4,7 +4,7 @@ from sqlalchemy import text from typing_extensions import override -from pgai.vectorizer.configuration import ( +from pgai.alembic.configuration import ( CharacterTextSplitterConfig, ChunkValueConfig, CreateVectorizerParams, @@ -23,8 +23,6 @@ class CreateVectorizerOp(MigrateOperation): - """Create a vectorizer for automatic embedding generation.""" - def __init__( self, source_table: str | None, @@ -68,19 +66,15 @@ def __init__( @classmethod def create_vectorizer(cls, operations: Operations, source_table: str, **kw: Any): - """Issue a CREATE VECTORIZER command.""" op = CreateVectorizerOp(source_table, **kw) return operations.invoke(op) @override def reverse(self) -> MigrateOperation: - """Creates the downgrade operation""" return DropVectorizerOp(None, True) class DropVectorizerOp(MigrateOperation): - """Drop a vectorizer and its associated objects.""" - def __init__(self, vectorizer_id: int | None, drop_all: bool): self.vectorizer_id = vectorizer_id self.drop_all = drop_all @@ -92,28 +86,22 @@ def drop_vectorizer( vectorizer_id: int | None = None, drop_all: bool = True, ): - """Issue a DROP VECTORIZER command.""" op = DropVectorizerOp(vectorizer_id, drop_all) return operations.invoke(op) @override def reverse(self) -> MigrateOperation: - """Creates the upgrade operation""" return CreateVectorizerOp(None) def create_vectorizer(operations: Operations, operation: CreateVectorizerOp): - """Implement CREATE VECTORIZER.""" params = operation.params operations.execute(params.to_sql()) def drop_vectorizer(operations: Operations, operation: DropVectorizerOp): - """Implement DROP VECTORIZER with cleanup of dependent objects.""" connection = operations.get_bind() vectorizer_id = operation.vectorizer_id - - # Finally drop the vectorizer itself connection.execute( text("SELECT ai.drop_vectorizer(:id, drop_all=>:drop_all)"), {"id": vectorizer_id, "drop_all": operation.drop_all}, diff --git a/projects/pgai/tests/vectorizer/extensions/fixtures/migrations/generic_vectorizer.py.template b/projects/pgai/tests/vectorizer/extensions/fixtures/migrations/generic_vectorizer.py.template index 526639b17..eac1c647e 100644 --- a/projects/pgai/tests/vectorizer/extensions/fixtures/migrations/generic_vectorizer.py.template +++ b/projects/pgai/tests/vectorizer/extensions/fixtures/migrations/generic_vectorizer.py.template @@ -7,7 +7,7 @@ Create Date: {create_date} from alembic import op from sqlalchemy import Column, Integer, String, Text from sqlalchemy import text -from pgai.vectorizer.configuration import ( +from pgai.alembic.configuration import ( OpenAIConfig, OllamaConfig, VoyageAIConfig, diff --git a/projects/pgai/tests/vectorizer/extensions/test_alembic.py b/projects/pgai/tests/vectorizer/extensions/test_alembic.py index d491dd314..6933aa6e7 100644 --- a/projects/pgai/tests/vectorizer/extensions/test_alembic.py +++ b/projects/pgai/tests/vectorizer/extensions/test_alembic.py @@ -7,7 +7,7 @@ If you add new configuration options, add a test for them here. Either in an existing test case or a new one. -The config classes for migrations lie in pgai/vectorizer/configuration.py +The config classes for migrations lie in pgai/alembic/configuration.py If the migration and the validation definitions overlap, it is possible to define a parent class in base.py see e.g. CharacterTextSplitterConfig as an example @@ -22,7 +22,7 @@ from alembic.config import Config from sqlalchemy import Engine, text -from pgai.vectorizer.configuration import ( +from pgai.alembic.configuration import ( CharacterTextSplitterConfig, HNSWIndexingConfig, OllamaConfig,