Ready to ship

SilasMarvin · SilasMarvin · commit 1d3fde69665b · 2024-08-21T15:14:09.000-07:00
diff --git a/LICENSE b/LICENSE
@@ -0,0 +1,20 @@
+Copyright (c) 2024 PostgresML Team
+
+Permission is hereby granted, free of charge, to any person obtaining
+a copy of this software and associated documentation files (the
+"Software"), to deal in the Software without restriction, including
+without limitation the rights to use, copy, modify, merge, publish,
+distribute, sublicense, and/or sell copies of the Software, and to
+permit persons to whom the Software is furnished to do so, subject to
+the following conditions:
+
+The above copyright notice and this permission notice shall be
+included in all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
+EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
+MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
+NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
+LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
+OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
+WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
diff --git a/README.md b/README.md
@@ -0,0 +1,76 @@
+# postgresml-django
+
+postgresml-django is a Python module that integrates PostgresML with Django ORM, enabling automatic in-database embedding of Django models. It simplifies the process of creating and searching vector embeddings for your text data.
+
+## Introduction
+
+This module provides a seamless way to:
+- Automatically generate in-databse embeddings for specified fields in your Django models
+- Perform vector similarity searches in-database
+
+## Installation
+
+1. Ensure you have [pgml](https://github.com/postgresml/postgresml) installed and configured in your database. The easiest way to do that is to sign up for a free serverless database at [postgresml.org](https://postgresml.org). You can also host it your self.
+
+2. Install the package using pip:
+
+   ```
+   pip install postgresml-django
+   ```
+
+You are ready to go!
+
+## Usage Examples
+
+### Example 1: Using intfloat/e5-small-v2
+
+This example demonstrates using the `intfloat/e5-small-v2` transformer, which has an embedding size of 384.
+
+```python
+from django.db import models
+from postgresml_django import VectorField, Embed
+
+class Document(Embed):
+    text = models.TextField()
+    text_embedding = VectorField(
+        field_to_embed="text",
+        dimensions=384,
+        transformer="intfloat/e5-small-v2"
+    )
+
+# Searching
+results = Document.vector_search("text_embedding", "some query to search against")
+```
+
+### Example 2: Using mixedbread-ai/mxbai-embed-large-v1
+
+This example shows how to use the `mixedbread-ai/mxbai-embed-large-v1` transformer, which has an embedding size of 512 and requires specific parameters for recall.
+
+```python
+from django.db import models
+from postgresml_django import VectorField, Embed
+
+class Article(Embed):
+    content = models.TextField()
+    content_embedding = VectorField(
+        field_to_embed="content",
+        dimensions=512,
+        transformer="mixedbread-ai/mxbai-embed-large-v1",
+        transformer_recall_parameters={
+            "query": "Represent this sentence for searching relevant passages: "
+        }
+    )
+
+# Searching
+results = Article.vector_search("content_embedding", "search query")
+```
+
+Note the differences between the two examples:
+1. The `dimensions` parameter is set to 384 for `intfloat/e5-small-v2` and 512 for `mixedbread-ai/mxbai-embed-large-v1`.
+2. The `mixedbread-ai/mxbai-embed-large-v1` transformer requires additional parameters for recall, which are specified in the `transformer_recall_parameters` argument.
+
+Both examples will automatically generate embeddings when instances are saved and allow for vector similarity searches using the `vector_search` method.
+
+## Contributing
+
+We welcome contributions to postgresml-django! Whether it's bug reports, feature requests, documentation improvements, or code contributions, your input is valuable to us. Feel free to open issues or submit pull requests on our GitHub repository.
diff --git a/pyproject.toml b/pyproject.toml
@@ -0,0 +1,28 @@
+[project]
+name = "postgresml-django"
+requires-python = ">=3.8"
+version = "0.1.0"
+description = "PostgresML Django integration that enables automatic embedding of specified fields."
+authors = [
+  {name = "PostgresML",  email = "team@postgresml.org"},
+]
+readme = "README.md"
+keywords = ["django","machine learning","vector databases","embeddings"]
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "License :: OSI Approved :: MIT License",
+    "Operating System :: OS Independent",
+]
+dependencies = [
+  "Django",
+  "pgvector"
+]
+
+[project.urls]
+Homepage = "https://postgresml.org"
+Repository = "https://github.com/postgresml/postgresml-django"
+Documentation = "https://github.com/postgresml/postgresml-django"
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
diff --git a/src/postgresml_django/__init__.py b/src/postgresml_django/__init__.py
@@ -0,0 +1 @@
+from .main import *
diff --git a/src/postgresml_django/main.py b/src/postgresml_django/main.py
@@ -0,0 +1,95 @@
+from django.db import models
+from django.db.models import Func, Value, F
+from django.db.models.functions import Cast
+import pgvector.django
+import json
+
+
+class GenerateEmbedding(Func):
+    function = "pgml.embed"
+    template = "%(function)s('%(transformer)s', %(expressions)s, '%(parameters)s')"
+    allowed_default = False
+
+    def __init__(self, expression, transformer, parameters={}):
+        self.transformer = transformer
+        self.parameters = parameters
+        super().__init__(expression)
+
+    def as_sql(self, compiler, connection, **extra_context):
+        extra_context["transformer"] = self.transformer
+        extra_context["parameters"] = json.dumps(self.parameters)
+        return super().as_sql(compiler, connection, **extra_context)
+
+
+class Embed(models.Model):
+    class Meta:
+        abstract = True
+
+    def save(self, *args, **kwargs):
+        update_fields = kwargs.get("update_fields")
+
+        # Check for fields to embed
+        for field in self._meta.get_fields():
+            if isinstance(field, VectorField):
+                if not hasattr(self, field.field_to_embed):
+                    raise AttributeError(
+                        f"Field to embed does not exist: `{field.field_to_embed}`"
+                    )
+
+                # Only embed if it's a new instance, full save, or this field is being updated
+                if not self.pk or update_fields is None or field.name in update_fields:
+                    value_to_embed = getattr(self, field.field_to_embed)
+                    setattr(
+                        self,
+                        field.name,
+                        GenerateEmbedding(
+                            Value(value_to_embed),
+                            field.transformer,
+                            field.transformer_store_parameters,
+                        ),
+                    )
+
+        super().save(*args, **kwargs)
+
+    @classmethod
+    def vector_search(
+        cls, field, query_text, distance_function=pgvector.django.CosineDistance
+    ):
+        # Get the fields
+        field_instance = getattr(cls._meta.model, field).field
+
+        # Generate an embedding for the text
+        query_embedding = GenerateEmbedding(
+            Value(query_text),
+            "intfloat/e5-small-v2",
+            field_instance.transformer_recall_parameters,
+        )
+
+        # Return the QuerySet
+        return cls.objects.annotate(
+            distance=distance_function(
+                F(field),
+                Cast(
+                    query_embedding,
+                    output_field=VectorField(dimensions=field_instance.dimensions),
+                ),
+            )
+        ).order_by("distance")
+
+
+class VectorField(pgvector.django.VectorField):
+    def __init__(
+        self,
+        field_to_embed=None,
+        dimensions=None,
+        transformer=None,
+        transformer_store_parameters={},
+        transformer_recall_parameters={},
+        *args,
+        **kwargs,
+    ):
+        self.field_to_embed = field_to_embed
+        self.transformer = transformer
+        self.transformer_store_parameters = transformer_store_parameters
+        self.transformer_recall_parameters = transformer_recall_parameters
+        super().__init__(dimensions=dimensions, *args, **kwargs)
