feat: add alembic operations for vectorizer #266

Askir · 2024-12-02T10:03:17Z

This PR adds native python operations to alembic so you don't have to write SQL to create vectorizers.

cevian

I gotta say I'm not convinced about the arguments for using a separate model than the models already in pgai/vectorizer or at least having both sets of models extend a base model. I think having 2 sets of models with similar params is really hard to maintain and quite a bit of code duplication. I'd like some more eyes on this tho. Can
James and/or Alejandro chime in here. In particular I'd like us to consider three designs:

simply extending the pydantic model we already have with optional fields that are present in either the stored json OR needed for the alembic stuff + having some kind of wrappers to create the config objects in alembic.
Factoring common data fields into base classes and using those as mixins. (kinda like the ApiKeyMixin now).
Maybe I'm just being stubborn and we should have separate models, like Jascha has them now.
leaving a few comments in but I think this is the big issue we need to resolve

cevian · 2024-12-04T19:54:08Z

projects/pgai/pgai/configuration.py

+        parts.append(")")
+        return "\n".join(parts)
+
+    def to_python(self, autogen_context: AutogenContext) -> str:


The to_python stuff is only needed if we do autogeneration right? So should really be in the 3rd PR?

cevian · 2024-12-04T19:57:20Z

projects/pgai/pgai/alembic/operations.py

+_operations_registered = False
+
+
+def register_operations():


is this also only needed for the autogen stuff?

no this is necessary to make the create and drop vectorizer functions available under alembic op context.
So you can call op.create_vectorizer

All the alembic operations are managed this way. You can read up about the register functions here. They are meant to be used as a decorator but then they only work if the module is imported. So I think asking the user to explicitly call register_operations, is less confusing.

cevian · 2024-12-04T19:58:55Z

projects/pgai/pgai/alembic/operations.py

+    @override
+    def reverse(self) -> MigrateOperation:
+        """Creates the upgrade operation"""
+        return CreateVectorizerOp(None)


what is the None argument for?

The reverse function is used to allow reversing migrations (upgrade->downgrade) and it seems to have to create a valid MIgrateOperation Object.
However to reverse a drop operation I'd have to fully load the current config and pass all the parameters back into the create operations (similar to what the autogenerate PR does).

To avoid this I made the CreateVectorizerOps source_table nullable, this doesn't create a valid operation since you can't call ai.create_vectorizer(NULL) but at least the alembic interface is happy.

cevian · 2024-12-04T20:47:01Z

projects/pgai/pgai/configuration.py

+    def to_python_arg(self) -> str: ...
+
+
+def format_python_arg(config_type: str, instance: Any) -> str:


can't we right a similar function for sql to make all this simpler to write? I mean all parameters can be passed in with {name} => {value} and the way to render the value is type dependent, which we can switch on. We can also skip any None's etc

JamesGuthrie · 2024-12-05T10:56:18Z

docs/python-integration.md

-        'content',
-        50,  -- chunk_size
-        10   -- chunk_overlap
+pgai provides native Alembic operations for managing vectorizers. For them to work you need to run `setup_alembic` in your env.py file. Which registers the pgai operations under the global op context:


What is setup_alembic? Is it something that people familiar with alembic know about?

I renamed this to register_operations now and explained it here: #266 (comment)
Hope that makes it a bit more clear what this is for.

… on first query

Askir · 2024-12-13T10:54:03Z

projects/pgai/pgai/alembic/operations.py

+    @override
+    def reverse(self) -> MigrateOperation:
+        """Creates the upgrade operation"""
+        return CreateVectorizerOp(None)


The reverse function is used to allow reversing migrations (upgrade->downgrade) and it seems to have to create a valid MIgrateOperation Object.
However to reverse a drop operation I'd have to fully load the current config and pass all the parameters back into the create operations (similar to what the autogenerate PR does).

To avoid this I made the CreateVectorizerOps source_table nullable, this doesn't create a valid operation since you can't call ai.create_vectorizer(NULL) but at least the alembic interface is happy.

Askir · 2024-12-13T10:56:39Z

projects/pgai/pgai/alembic/operations.py

+_operations_registered = False
+
+
+def register_operations():


no this is necessary to make the create and drop vectorizer functions available under alembic op context.
So you can call op.create_vectorizer

All the alembic operations are managed this way. You can read up about the register functions here. They are meant to be used as a decorator but then they only work if the module is imported. So I think asking the user to explicitly call register_operations, is less confusing.

Askir · 2024-12-13T11:01:32Z

projects/pgai/tests/vectorizer/extensions/fixtures/migrations/002_create_vectorizer.py.template

+        chunking=CharacterTextSplitterConfig(
+            chunk_column='content',
+            chunk_size=800,
+            chunk_overlap=400,
+            separator='.',
+            is_separator_regex=False
+        ),


@cevian This is a bit of an ugly side effect of having merged the parent classes now.
The pydantic models that map the db config have these fields as non optional because the extension code provides the defaults.

But the input function should have these as optional for the same reason (if they are null the extension will provide e.g. chunk_size=800.

So now in the python code you have to unnecessarily always provide all the arguments (true for some of the other configs as well).

Alternatively I could give these the same defaults as the extension code provides, but then we have to again maintain the defaults in 2 places.

This was previously not an issue as I could just make the input arguments nullable while leaving the load from db models required. I can override this of course in the child classes, but then again this breaks the purpose of sharing a parent class in the first place.

Askir force-pushed the jascha/add-alembic-migration-ops branch from c899380 to fd9f1bc Compare December 2, 2024 10:08

Askir mentioned this pull request Dec 2, 2024

feat: SQLAlchemy and alembic integration #208

Closed

Askir force-pushed the jascha/add-alembic-migration-ops branch from fd9f1bc to 6f5ff59 Compare December 3, 2024 13:37

Askir marked this pull request as ready for review December 3, 2024 23:16

Askir requested a review from a team as a code owner December 3, 2024 23:16

Askir force-pushed the jascha/add-vectorizer-field branch from 8742af8 to 36cf4d9 Compare December 4, 2024 13:13

Askir force-pushed the jascha/add-alembic-migration-ops branch from 6f5ff59 to 525ab5b Compare December 4, 2024 13:20

cevian requested changes Dec 4, 2024

View reviewed changes

JamesGuthrie reviewed Dec 5, 2024

View reviewed changes

Askir force-pushed the jascha/add-vectorizer-field branch 3 times, most recently from e342f0b to c108a87 Compare December 7, 2024 06:34

Askir added 13 commits December 12, 2024 12:36

feat: add sqlalchemy vectorizer field

a011739

chore: add installing extras to ci

dbd39ad

chore: simplify interface, add simple docs

e2afe66

feat: allow arbitrary primary keys on parent model

4637367

docs: update docs with simplified vectorizer field

a0414df

chore: rename VectorizerField to Vectorizer

86a952d

chore: update alembic exclusion mechanism

67f5209

docs: update docs with review comments

11de6d3

chore: align automatic table name with create_vectorizer

b6096e0

chore: add option for any relationship properties

9cc46ac

chore: setup class event based rather than lazy so relationship works…

800d160

… on first query

chore: update to embedding_relationship

cff7528

chore: allow api keys in lib tests

da7b72b

Askir force-pushed the jascha/add-vectorizer-field branch 5 times, most recently from 6bb0414 to 4c8e1a0 Compare December 12, 2024 12:55

Askir force-pushed the jascha/add-vectorizer-field branch from 4c8e1a0 to 3b47afc Compare December 12, 2024 13:33

chore: refactor tests add vcr mocks

8fe145e

Askir force-pushed the jascha/add-vectorizer-field branch from 3b47afc to 8fe145e Compare December 12, 2024 13:46

Askir added 2 commits December 12, 2024 17:42

feat: add alembic operations for vectorizer

ff28a39

chore: cleanup set up of operations

5e76cf9

Askir force-pushed the jascha/add-alembic-migration-ops branch from 525ab5b to 5e76cf9 Compare December 12, 2024 16:44

Askir added 2 commits December 13, 2024 11:41

chore: add shared base class

5b11d7e

docs: update docs

0d96b9a

Askir commented Dec 13, 2024

View reviewed changes

Askir added 2 commits December 13, 2024 14:38

chore: unify sql generation

edec725

chore: add more test cases

c619acd

Askir force-pushed the jascha/add-vectorizer-field branch from 8fe145e to 882f91e Compare December 19, 2024 11:40

Base automatically changed from jascha/add-vectorizer-field to main December 19, 2024 12:32

chore: simplify code and tests a bit

4792820

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat: add alembic operations for vectorizer #266

feat: add alembic operations for vectorizer #266

Askir commented Dec 2, 2024 •

edited

Loading

cevian left a comment

cevian Dec 4, 2024

cevian Dec 4, 2024

Askir Dec 13, 2024

cevian Dec 4, 2024

Askir Dec 13, 2024

cevian Dec 4, 2024

JamesGuthrie Dec 5, 2024

Askir Dec 13, 2024

Askir Dec 13, 2024

Askir Dec 13, 2024

Askir Dec 13, 2024 •

edited

Loading

		def to_python_arg(self) -> str: ...


		def format_python_arg(config_type: str, instance: Any) -> str:

feat: add alembic operations for vectorizer #266

Are you sure you want to change the base?

feat: add alembic operations for vectorizer #266

Conversation

Askir commented Dec 2, 2024 • edited Loading

cevian left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Askir Dec 13, 2024 • edited Loading

Choose a reason for hiding this comment

Askir commented Dec 2, 2024 •

edited

Loading

Askir Dec 13, 2024 •

edited

Loading