chore: move packages around, add sqlalchemy as optional extra

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

Hide file tree

Showing 13 changed files with 28 additions and 21 deletions.
diff --git a/docs/python-integration.md b/docs/python-integration.md
@@ -240,17 +240,18 @@ Configure your env.py to support vectorizer autogeneration:
 
 ```python
 from alembic import context
-from pgai.extensions.alembic.autogenerate import compare_vectorizers
-from pgai.extensions.alembic.operations import CreateVectorizerOp, DropVectorizerOp
+from pgai.alembic import compare_vectorizers
+from pgai.alembic import CreateVectorizerOp, DropVectorizerOp
+
 
 # Make sure your env.py includes:
 def run_migrations_online():
     with connectable.connect() as connection:
         context.configure(
             connection=connection,
             target_metadata=target_metadata,
-            include_object=lambda obj, name, type_, reflected, compare_to: 
-                not obj.info.get("pgai_managed", False)
+            include_object=lambda obj, name, type_, reflected, compare_to:
+            not obj.info.get("pgai_managed", False)
         )
 ```
 

diff --git a/projects/pgai/pgai/alembic/__init__.py b/projects/pgai/pgai/alembic/__init__.py
@@ -0,0 +1 @@
+__all__ = ["CreateVectorizerOp", "DropVectorizerOp"]
diff --git a/...i/pgai/extensions/alembic/autogenerate.py → projects/pgai/pgai/alembic/autogenerate.py b/...i/pgai/extensions/alembic/autogenerate.py → projects/pgai/pgai/alembic/autogenerate.py
@@ -5,7 +5,7 @@
 from alembic.operations.ops import UpgradeOps
 from sqlalchemy import text
 
-from pgai.extensions.alembic.operations import (
+from pgai.alembic import (
     ChunkingConfig,
     CreateVectorizerOp,
     DiskANNIndexingConfig,

diff --git a/...gai/pgai/extensions/alembic/operations.py → projects/pgai/pgai/alembic/operations.py b/...gai/pgai/extensions/alembic/operations.py → projects/pgai/pgai/alembic/operations.py
diff --git a/projects/pgai/pgai/extensions/__init__.py → projects/pgai/pgai/configuration.py b/projects/pgai/pgai/extensions/__init__.py → projects/pgai/pgai/configuration.py
diff --git a/projects/pgai/pgai/extensions/alembic/__init__.py b/projects/pgai/pgai/extensions/alembic/__init__.py
diff --git a/...ai/pgai/extensions/sqlalchemy/__init__.py → projects/pgai/pgai/sqlalchemy/__init__.py b/...ai/pgai/extensions/sqlalchemy/__init__.py → projects/pgai/pgai/sqlalchemy/__init__.py
diff --git a/projects/pgai/pyproject.toml b/projects/pgai/pyproject.toml
@@ -23,9 +23,6 @@ dependencies = [
     "typing_extensions>=4.0,<5.0",
     "datadog_lambda>=6.9,<7.0",
     "pytimeparse>=1.1,<2.0",
-    "sqlalchemy>=2.0.36",
-    "psycopg2>=2.9.10",
-    "alembic>=1.14.0",
 ]
 classifiers = [
     "License :: OSI Approved :: PostgreSQL License",
@@ -44,6 +41,13 @@ classifiers = [
     "Operating System :: POSIX :: Linux",
 ]
 
+[project.optional-dependencies]
+sqlalchemy=[
+    "sqlalchemy>=2.0.36",
+    "psycopg2>=2.9.10",
+    "alembic>=1.14.0",
+]
+
 [project.urls]
 Homepage = "https://github.com/timescale/pgai"
 Repository = "https://github.com/timescale/pgai"

diff --git a/projects/pgai/tests/vectorizer/extensions/test_sqlalchemy.py b/projects/pgai/tests/vectorizer/extensions/test_sqlalchemy.py
@@ -5,7 +5,7 @@
 from testcontainers.postgres import PostgresContainer  # type: ignore
 
 from pgai.cli import vectorizer_worker
-from pgai.extensions.alembic.operations import (
+from pgai.alembic import (
     ChunkingConfig,
     EmbeddingConfig,
 )

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,7 +6,7 @@
 from testcontainers.postgres import PostgresContainer  # type: ignore
 
 from pgai.cli import vectorizer_worker
-from pgai.extensions.alembic.operations import (
+from pgai.alembic import (
     EmbeddingConfig,
 )
 from pgai.extensions.sqlalchemy import VectorizerField

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,7 +6,7 @@
 from testcontainers.postgres import PostgresContainer  # type: ignore
 
 from pgai.cli import vectorizer_worker
-from pgai.extensions.alembic.operations import (
+from pgai.alembic import (
     ChunkingConfig,
     EmbeddingConfig,
 )

diff --git a/projects/pgai/tests/vectorizer/extensions/test_sqlalchemy_relationship.py b/projects/pgai/tests/vectorizer/extensions/test_sqlalchemy_relationship.py
@@ -6,7 +6,7 @@
 from testcontainers.postgres import PostgresContainer  # type: ignore
 
 from pgai.cli import vectorizer_worker
-from pgai.extensions.alembic.operations import (
+from pgai.alembic import (
     ChunkingConfig,
     EmbeddingConfig,
 )

diff --git a/projects/pgai/uv.lock b/projects/pgai/uv.lock
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1 @@
		__all__ = ["CreateVectorizerOp", "DropVectorizerOp"]