chore: change package structure

timescale · Nov 22, 2024 · 84a27f4 · 84a27f4
1 parent b5e8d64
commit 84a27f4
Show file tree

Hide file tree

Showing 12 changed files with 32 additions and 28 deletions.
diff --git a/projects/pgai/pgai/alembic/__init__.py b/projects/pgai/pgai/alembic/__init__.py
@@ -1 +1,5 @@
-__all__ = ["CreateVectorizerOp", "DropVectorizerOp"]
+from pgai.alembic.operations import CreateVectorizerOp, DropVectorizerOp
+from pgai.alembic.autogenerate import compare_vectorizers
+
+
+__all__ = ["CreateVectorizerOp", "DropVectorizerOp", "compare_vectorizers"]
diff --git a/projects/pgai/pgai/alembic/autogenerate.py b/projects/pgai/pgai/alembic/autogenerate.py
@@ -5,16 +5,16 @@
 from alembic.operations.ops import UpgradeOps
 from sqlalchemy import text
 
-from pgai.alembic import (
+from pgai.configuration import (
     ChunkingConfig,
-    CreateVectorizerOp,
     DiskANNIndexingConfig,
-    DropVectorizerOp,
     EmbeddingConfig,
     ProcessingConfig,
     SchedulingConfig,
 )
 
+from pgai.alembic.operations import CreateVectorizerOp, DropVectorizerOp
+
 
 @comparators.dispatch_for("schema")
 def compare_vectorizers(

diff --git a/projects/pgai/pgai/alembic/operations.py b/projects/pgai/pgai/alembic/operations.py
@@ -6,7 +6,7 @@
 from sqlalchemy import text
 from typing_extensions import override
 
-from pgai.extensions import (
+from pgai.configuration import (
     ChunkingConfig,
     DiskANNIndexingConfig,
     EmbeddingConfig,
@@ -188,12 +188,12 @@ def drop_vectorizer(operations: Operations, operation: DropVectorizerOp):
 def render_create_vectorizer(autogen_context: AutogenContext, op: CreateVectorizerOp):
     """Render a CREATE VECTORIZER operation."""
     template_context = {
-        "EmbeddingConfig": "from pgai.extensions import EmbeddingConfig",
-        "ChunkingConfig": "from pgai.extensions import ChunkingConfig",
-        "DiskANNIndexingConfig": "from pgai.extensions import DiskANNIndexingConfig",
-        "HNSWIndexingConfig": "from pgai.extensions import HNSWIndexingConfig",
-        "SchedulingConfig": "from pgai.extensions import SchedulingConfig",
-        "ProcessingConfig": "from pgai.extensions import ProcessingConfig",
+        "EmbeddingConfig": "from pgai.configuration import EmbeddingConfig",
+        "ChunkingConfig": "from pgai.configuration import ChunkingConfig",
+        "DiskANNIndexingConfig": "from pgai.configuration import DiskANNIndexingConfig",
+        "HNSWIndexingConfig": "from pgai.configuration import HNSWIndexingConfig",
+        "SchedulingConfig": "from pgai.configuration import SchedulingConfig",
+        "ProcessingConfig": "from pgai.configuration import ProcessingConfig",
     }
 
     for import_str in template_context.values():

diff --git a/projects/pgai/pgai/sqlalchemy/__init__.py b/projects/pgai/pgai/sqlalchemy/__init__.py
@@ -5,7 +5,7 @@
 from sqlalchemy import ForeignKey, Integer, MetaData, Text
 from sqlalchemy.orm import DeclarativeBase, Mapped, backref, mapped_column, relationship
 
-from pgai.extensions import (
+from pgai.configuration import (
     ChunkingConfig,
     DiskANNIndexingConfig,
     EmbeddingConfig,

diff --git a/projects/pgai/tests/vectorizer/extensions/fixtures/alembic/autogen_env.py.template b/projects/pgai/tests/vectorizer/extensions/fixtures/alembic/autogen_env.py.template
@@ -5,9 +5,9 @@ from logging.config import fileConfig
 from alembic import context
 from sqlalchemy import engine_from_config, pool
 
-from pgai.extensions.alembic.autogenerate import compare_vectorizers
-from pgai.extensions.alembic.operations import CreateVectorizerOp, DropVectorizerOp
-from pgai.extensions.alembic.operations import render_create_vectorizer, render_drop_vectorizer
+from pgai.alembic.autogenerate import compare_vectorizers
+from pgai.alembic.operations import CreateVectorizerOp, DropVectorizerOp
+from pgai.alembic.operations import render_create_vectorizer, render_drop_vectorizer
 
 sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 

diff --git a/...ts/pgai/tests/vectorizer/extensions/fixtures/migrations/002_create_vectorizer.py.template b/...ts/pgai/tests/vectorizer/extensions/fixtures/migrations/002_create_vectorizer.py.template
@@ -5,8 +5,8 @@ Revises: {revises}
 Create Date: {create_date}
 """
 from alembic import op
-from pgai.extensions.alembic import CreateVectorizerOp
-from pgai.extensions import (
+from pgai.alembic import CreateVectorizerOp
+from pgai.configuration import (
     EmbeddingConfig,
     ChunkingConfig
 )

diff --git a/projects/pgai/tests/vectorizer/extensions/fixtures/models/blog_post.py.template b/projects/pgai/tests/vectorizer/extensions/fixtures/models/blog_post.py.template
@@ -1,6 +1,6 @@
 from sqlalchemy.orm import declarative_base
 from sqlalchemy import Column, Integer, Text
-from pgai.extensions.sqlalchemy import (
+from pgai.sqlalchemy import (
     VectorizerField,
     EmbeddingConfig,
     ChunkingConfig

diff --git a/projects/pgai/tests/vectorizer/extensions/fixtures/models/blog_post_all_fields.py.template b/projects/pgai/tests/vectorizer/extensions/fixtures/models/blog_post_all_fields.py.template
@@ -1,8 +1,8 @@
 from sqlalchemy.orm import declarative_base
 from sqlalchemy import Column, Integer, Text
 
-from pgai.extensions import DiskANNIndexingConfig, SchedulingConfig, ProcessingConfig
-from pgai.extensions.sqlalchemy import (
+from pgai.configuration import DiskANNIndexingConfig, SchedulingConfig, ProcessingConfig
+from pgai.sqlalchemy import (
     VectorizerField,
     EmbeddingConfig,
     ChunkingConfig

diff --git a/projects/pgai/tests/vectorizer/extensions/test_sqlalchemy.py b/projects/pgai/tests/vectorizer/extensions/test_sqlalchemy.py
@@ -5,11 +5,11 @@
 from testcontainers.postgres import PostgresContainer  # type: ignore
 
 from pgai.cli import vectorizer_worker
-from pgai.alembic import (
+from pgai.configuration import (
     ChunkingConfig,
     EmbeddingConfig,
 )
-from pgai.extensions.sqlalchemy import VectorizerField
+from pgai.sqlalchemy import VectorizerField
 
 
 def test_sqlalchemy(postgres_container: PostgresContainer, initialized_engine: Engine):

diff --git a/projects/pgai/tests/vectorizer/extensions/test_sqlalchemy_large_embeddings.py b/projects/pgai/tests/vectorizer/extensions/test_sqlalchemy_large_embeddings.py
@@ -6,10 +6,10 @@
 from testcontainers.postgres import PostgresContainer  # type: ignore
 
 from pgai.cli import vectorizer_worker
-from pgai.alembic import (
+from pgai.configuration import (
     EmbeddingConfig,
 )
-from pgai.extensions.sqlalchemy import VectorizerField
+from pgai.sqlalchemy import VectorizerField
 
 
 class Base(DeclarativeBase):

diff --git a/projects/pgai/tests/vectorizer/extensions/test_sqlalchemy_no_relationship.py b/projects/pgai/tests/vectorizer/extensions/test_sqlalchemy_no_relationship.py
@@ -6,11 +6,11 @@
 from testcontainers.postgres import PostgresContainer  # type: ignore
 
 from pgai.cli import vectorizer_worker
-from pgai.alembic import (
+from pgai.configuration import (
     ChunkingConfig,
     EmbeddingConfig,
 )
-from pgai.extensions.sqlalchemy import VectorizerField
+from pgai.sqlalchemy import VectorizerField
 
 
 class Base(DeclarativeBase):

diff --git a/projects/pgai/tests/vectorizer/extensions/test_sqlalchemy_relationship.py b/projects/pgai/tests/vectorizer/extensions/test_sqlalchemy_relationship.py
@@ -6,11 +6,11 @@
 from testcontainers.postgres import PostgresContainer  # type: ignore
 
 from pgai.cli import vectorizer_worker
-from pgai.alembic import (
+from pgai.configuration import (
     ChunkingConfig,
     EmbeddingConfig,
 )
-from pgai.extensions.sqlalchemy import EmbeddingModel, VectorizerField
+from pgai.sqlalchemy import EmbeddingModel, VectorizerField
 
 
 class Base(DeclarativeBase):