corvic-engine 0.3.0rc76__tar.gz → 0.3.0rc78__tar.gz

{corvic_engine-0.3.0rc76 → corvic_engine-0.3.0rc78}/Cargo.lock

@@ -16,7 +16,7 @@ checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd"
 
 [[package]]
 name = "corvic-engine"
-version = "0.3.0-rc.76"
+version = "0.3.0-rc.78"
 dependencies = [
  "numpy",
  "pyo3",
@@ -57,7 +57,7 @@ checksum = "48c757948c5ede0e46177b7add2e67155f70e33c07fea8284df6576da70b3719"
 
 [[package]]
 name = "engine-walk"
-version = "0.3.0-rc.76"
+version = "0.3.0-rc.78"
 dependencies = [
  "version_check",
 ]
@@ -82,9 +82,9 @@ checksum = "d750af042f7ef4f724306de029d18836c26c1765a54a6a3f094cbd23a7267ffa"
 
 [[package]]
 name = "matrixmultiply"
-version = "0.3.9"
+version = "0.3.10"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "9380b911e3e96d10c1f415da0876389aaf1b56759054eeb0de7df940c456ba1a"
+checksum = "a06de3016e9fae57a36fd14dba131fccf49f74b40b7fbdb472f96e361ec71a08"
 dependencies = [
  "autocfg",
  "rawpointer",
@@ -293,9 +293,9 @@ checksum = "357703d41365b4b27c590e3ed91eabb1b663f07c4c084095e60cbed4362dff0d"
 
 [[package]]
 name = "syn"
-version = "2.0.100"
+version = "2.0.101"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b09a44accad81e1ba1cd74a32461ba89dee89095ba17b32f5d03683b1b1fc2a0"
+checksum = "8ce2b7fc941b3a24138a0a7cf8e858bfc6a992e7978a068a5c760deb0ed43caf"
 dependencies = [
  "proc-macro2",
  "quote",

{corvic_engine-0.3.0rc76 → corvic_engine-0.3.0rc78}/Cargo.toml

@@ -3,7 +3,7 @@ resolver = "2"
 members = ["python"]
 
 [workspace.package]
-version = "0.3.0-rc.76"
+version = "0.3.0-rc.78"
 edition = "2021"
 publish = false
 

{corvic_engine-0.3.0rc76 → corvic_engine-0.3.0rc78}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: corvic-engine
-Version: 0.3.0rc76
+Version: 0.3.0rc78
 Classifier: Environment :: Console
 Classifier: License :: Other/Proprietary License
 Classifier: Programming Language :: Python :: Implementation :: CPython

{corvic_engine-0.3.0rc76 → corvic_engine-0.3.0rc78}/pyproject.toml

@@ -18,7 +18,7 @@ authors = [
 ]
 readme = "README.md"
 description = "Seamless embedding generation and retrieval."
-version = "0.3.0-rc.76" # sync this with tool.poetry.version below
+version = "0.3.0-rc.78" # sync this with tool.poetry.version below
 
 # TODO(ddn): Pull dependencies from tool.poetry.dependencies. We use maturin as
 # the build system but poetry for the development environment.
@@ -72,7 +72,7 @@ packages = [
 ]
 description = ""
 authors = []
-version = "0.3.0-rc.76"
+version = "0.3.0-rc.78"
 
 [tool.poetry.scripts]
 check = "corvic_check.cli:main"
@@ -120,6 +120,7 @@ import-linter = "^2.0"
 maturin = "^1.4.0"
 pdoc = "^14.5.1"
 pre-commit = "^4.0.0"
+psycopg = { extras = ["binary"], version = "^3.2" }
 pyright = "^1.1.387"
 pytest = "^8.1.1"
 pytest-asyncio = "^0.26"
@@ -403,6 +404,7 @@ root_packages = [
   "corvic.embed",
   "corvic.embedding_metric",
   "corvic.engine",
+  "corvic.eorm",
   "corvic.model",
   "corvic.op_graph",
   "corvic.orm",

{corvic_engine-0.3.0rc76/python/corvic/orm → corvic_engine-0.3.0rc78/python/corvic/eorm}/__init__.py

@@ -1,52 +1,14 @@
-"""Data model definitions; backed by an RDBMS."""
+"""Engine Object-Relational Mappings."""
 
 from __future__ import annotations
 
 from datetime import datetime
+from typing import TypeAlias
 
 import sqlalchemy as sa
 from sqlalchemy import orm as sa_orm
 
-from corvic.orm._proto_columns import ProtoMessageDecorator
-from corvic.orm.base import Base, OrgBase
-from corvic.orm.errors import (
-    DeletedObjectError,
-    InvalidORMIdentifierError,
-    RequestedObjectsForNobodyError,
-)
-from corvic.orm.ids import (
-    AgentID,
-    AgentMessageID,
-    BaseID,
-    BaseIDFromInt,
-    CompletionModelID,
-    FeatureViewID,
-    FeatureViewSourceID,
-    IntIDDecorator,
-    MessageEntryID,
-    OrgID,
-    PipelineID,
-    ResourceID,
-    RoomID,
-    SourceID,
-    SpaceID,
-    SpaceParametersID,
-    SpaceRunID,
-    UserMessageID,
-)
-from corvic.orm.keys import (
-    INT_PK_TYPE,
-    ForeignKey,
-    primary_key_foreign_column,
-    primary_key_identity_column,
-    primary_key_uuid_column,
-)
-from corvic.orm.mixins import (
-    BelongsToOrgMixin,
-    Session,
-    SoftDeleteMixin,
-    live_unique_constraint,
-)
+from corvic import orm
 from corvic_generated.orm.v1 import (
     common_pb2,
     completion_model_pb2,
@@ -57,27 +19,87 @@ from corvic_generated.orm.v1 import (
 )
 from corvic_generated.status.v1 import event_pb2
 
-# NOTE: The only safe use of "sa_orm.relationship" uses the args:
-# `viewonly=True` and `init=False`. Writes quickly become
-# a complex mess when implementers of commit need to reason about
-# which sub-object should be updated.
-#
-# Rather, classes in corvic.model define their own commit protocols,
-# and if sub-orm-model updates are required they are explicit.
+OrgID: TypeAlias = orm.OrgID
+
+
+class RoomID(orm.BaseIDFromInt):
+    """A unique identifier for a room."""
+
+
+class ResourceID(orm.BaseIDFromInt):
+    """A unique identifier for a resource."""
+
+
+class SourceID(orm.BaseIDFromInt):
+    """A unique identifier for a source."""
+
+
+class PipelineID(orm.BaseIDFromInt):
+    """A unique identifier for a pipeline."""
+
 
+class FeatureViewID(orm.BaseIDFromInt):
+    """A unique identifier for a feature view."""
 
-class Org(SoftDeleteMixin, OrgBase, kw_only=True):
-    """An organization it a top level grouping of resources."""
 
+class FeatureViewSourceID(orm.BaseIDFromInt):
+    """A unique identifier for a source in a feature view."""
 
-class Room(BelongsToOrgMixin, SoftDeleteMixin, Base, kw_only=True):
+
+class SpaceID(orm.BaseIDFromInt):
+    """A unique identifier for a space."""
+
+
+class CompletionModelID(orm.BaseIDFromInt):
+    """A unique identifier for a completion model."""
+
+
+orm.Base.registry.update_type_annotation_map(
+    {
+        # proto  message column types
+        common_pb2.BlobUrlList: orm.ProtoMessageDecorator(common_pb2.BlobUrlList()),
+        feature_view_pb2.FeatureViewOutput: orm.ProtoMessageDecorator(
+            feature_view_pb2.FeatureViewOutput()
+        ),
+        common_pb2.AgentMessageMetadata: orm.ProtoMessageDecorator(
+            common_pb2.AgentMessageMetadata()
+        ),
+        space_pb2.SpaceParameters: orm.ProtoMessageDecorator(
+            space_pb2.SpaceParameters()
+        ),
+        table_pb2.TableComputeOp: orm.ProtoMessageDecorator(table_pb2.TableComputeOp()),
+        table_pb2.NamedTables: orm.ProtoMessageDecorator(table_pb2.NamedTables()),
+        common_pb2.RetrievedEntities: orm.ProtoMessageDecorator(
+            common_pb2.RetrievedEntities()
+        ),
+        pipeline_pb2.PipelineTransformation: orm.ProtoMessageDecorator(
+            pipeline_pb2.PipelineTransformation()
+        ),
+        event_pb2.Event: orm.ProtoMessageDecorator(event_pb2.Event()),
+        completion_model_pb2.CompletionModelParameters: orm.ProtoMessageDecorator(
+            completion_model_pb2.CompletionModelParameters()
+        ),
+        # ID types
+        RoomID: orm.IntIDDecorator(RoomID()),
+        ResourceID: orm.IntIDDecorator(ResourceID()),
+        SourceID: orm.IntIDDecorator(SourceID()),
+        PipelineID: orm.IntIDDecorator(PipelineID()),
+        FeatureViewID: orm.IntIDDecorator(FeatureViewID()),
+        FeatureViewSourceID: orm.IntIDDecorator(FeatureViewSourceID()),
+        SpaceID: orm.IntIDDecorator(SpaceID()),
+        CompletionModelID: orm.IntIDDecorator(CompletionModelID()),
+    }
+)
+
+
+class Room(orm.BelongsToOrgMixin, orm.SoftDeleteMixin, orm.Base, kw_only=True):
     """A Room is a logical collection of Documents."""
 
     __tablename__ = "room"
-    __table_args__ = (live_unique_constraint("name", "org_id"),)
+    __table_args__ = (orm.live_unique_constraint("name", "org_id"),)
 
     name: sa_orm.Mapped[str] = sa_orm.mapped_column(sa.Text, default=None)
-    id: sa_orm.Mapped[RoomID | None] = primary_key_identity_column()
+    id: sa_orm.Mapped[RoomID | None] = orm.primary_key_identity_column()
 
     @property
     def room_key(self):
@@ -86,27 +108,29 @@ class Room(BelongsToOrgMixin, SoftDeleteMixin, Base, kw_only=True):
 
 class BelongsToRoomMixin(sa_orm.MappedAsDataclass):
     room_id: sa_orm.Mapped[RoomID | None] = sa_orm.mapped_column(
-        ForeignKey(Room).make(ondelete="CASCADE"),
+        orm.ForeignKey(Room).make(ondelete="CASCADE"),
         nullable=True,
         default=None,
     )
 
 
-class DefaultObjects(Base, kw_only=True):
+class DefaultObjects(orm.Base, kw_only=True):
     """Holds the identifiers for default objects."""
 
     __tablename__ = "default_objects"
     default_org: sa_orm.Mapped[OrgID | None] = sa_orm.mapped_column(
-        ForeignKey(Org).make(ondelete="CASCADE"),
+        orm.ForeignKey(orm.Org).make(ondelete="CASCADE"),
         nullable=False,
     )
     default_room: sa_orm.Mapped[RoomID | None] = sa_orm.mapped_column(
-        ForeignKey(Room).make(ondelete="CASCADE"), nullable=True, default=None
+        orm.ForeignKey(Room).make(ondelete="CASCADE"), nullable=True, default=None
+    )
+    version: sa_orm.Mapped[int | None] = orm.primary_key_identity_column(
+        type_=orm.INT_PK_TYPE
     )
-    version: sa_orm.Mapped[int | None] = primary_key_identity_column(type_=INT_PK_TYPE)
 
 
-class Resource(BelongsToOrgMixin, BelongsToRoomMixin, Base, kw_only=True):
+class Resource(orm.BelongsToOrgMixin, BelongsToRoomMixin, orm.Base, kw_only=True):
     """A Resource is a reference to some durably stored file.
 
     E.g., a document could be a PDF file, an image, or a text transcript of a
@@ -122,7 +146,7 @@ class Resource(BelongsToOrgMixin, BelongsToRoomMixin, Base, kw_only=True):
     size: sa_orm.Mapped[int] = sa_orm.mapped_column(nullable=True)
     original_path: sa_orm.Mapped[str] = sa_orm.mapped_column(nullable=True)
     description: sa_orm.Mapped[str] = sa_orm.mapped_column(nullable=True)
-    id: sa_orm.Mapped[ResourceID | None] = primary_key_identity_column()
+    id: sa_orm.Mapped[ResourceID | None] = orm.primary_key_identity_column()
     latest_event: sa_orm.Mapped[event_pb2.Event | None] = sa_orm.mapped_column(
         default=None, nullable=True
     )
@@ -134,7 +158,7 @@ class Resource(BelongsToOrgMixin, BelongsToRoomMixin, Base, kw_only=True):
     )
 
 
-class Source(BelongsToOrgMixin, BelongsToRoomMixin, Base, kw_only=True):
+class Source(orm.BelongsToOrgMixin, BelongsToRoomMixin, orm.Base, kw_only=True):
     """A source."""
 
     __tablename__ = "source"
@@ -143,7 +167,7 @@ class Source(BelongsToOrgMixin, BelongsToRoomMixin, Base, kw_only=True):
     name: sa_orm.Mapped[str] = sa_orm.mapped_column(sa.Text)
     # protobuf describing the operations required to construct a table
     table_op_graph: sa_orm.Mapped[table_pb2.TableComputeOp] = sa_orm.mapped_column()
-    id: sa_orm.Mapped[SourceID | None] = primary_key_identity_column()
+    id: sa_orm.Mapped[SourceID | None] = orm.primary_key_identity_column()
 
     source_files: sa_orm.Mapped[common_pb2.BlobUrlList | None] = sa_orm.mapped_column(
         default=None
@@ -157,7 +181,7 @@ class Source(BelongsToOrgMixin, BelongsToRoomMixin, Base, kw_only=True):
         return self.name
 
 
-class Pipeline(BelongsToOrgMixin, BelongsToRoomMixin, Base, kw_only=True):
+class Pipeline(orm.BelongsToOrgMixin, BelongsToRoomMixin, orm.Base, kw_only=True):
     """A resource to source pipeline."""
 
     __tablename__ = "pipeline"
@@ -168,7 +192,7 @@ class Pipeline(BelongsToOrgMixin, BelongsToRoomMixin, Base, kw_only=True):
     )
     name: sa_orm.Mapped[str] = sa_orm.mapped_column()
     description: sa_orm.Mapped[str | None] = sa_orm.mapped_column()
-    id: sa_orm.Mapped[PipelineID | None] = primary_key_identity_column()
+    id: sa_orm.Mapped[PipelineID | None] = orm.primary_key_identity_column()
 
     outputs: sa_orm.Mapped[list[PipelineOutput]] = sa_orm.relationship(
         viewonly=True,
@@ -177,7 +201,7 @@ class Pipeline(BelongsToOrgMixin, BelongsToRoomMixin, Base, kw_only=True):
     )
 
 
-class PipelineInput(BelongsToOrgMixin, BelongsToRoomMixin, Base, kw_only=True):
+class PipelineInput(orm.BelongsToOrgMixin, BelongsToRoomMixin, orm.Base, kw_only=True):
     """Pipeline input resources."""
 
     __tablename__ = "pipeline_input"
@@ -187,15 +211,15 @@ class PipelineInput(BelongsToOrgMixin, BelongsToRoomMixin, Base, kw_only=True):
     name: sa_orm.Mapped[str]
     """A name the pipeline uses to refer to this input."""
 
-    pipeline_id: sa_orm.Mapped[PipelineID] = primary_key_foreign_column(
-        ForeignKey(Pipeline).make(ondelete="CASCADE")
+    pipeline_id: sa_orm.Mapped[PipelineID] = orm.primary_key_foreign_column(
+        orm.ForeignKey(Pipeline).make(ondelete="CASCADE")
     )
-    resource_id: sa_orm.Mapped[ResourceID] = primary_key_foreign_column(
-        ForeignKey(Resource).make(ondelete="CASCADE")
+    resource_id: sa_orm.Mapped[ResourceID] = orm.primary_key_foreign_column(
+        orm.ForeignKey(Resource).make(ondelete="CASCADE")
     )
 
 
-class PipelineOutput(BelongsToOrgMixin, BelongsToRoomMixin, Base, kw_only=True):
+class PipelineOutput(orm.BelongsToOrgMixin, BelongsToRoomMixin, orm.Base, kw_only=True):
     """Objects for tracking pipeline output sources."""
 
     __tablename__ = "pipeline_output"
@@ -205,23 +229,27 @@ class PipelineOutput(BelongsToOrgMixin, BelongsToRoomMixin, Base, kw_only=True):
     name: sa_orm.Mapped[str]
     """A name the pipeline uses to refer to this output."""
 
-    pipeline_id: sa_orm.Mapped[PipelineID] = primary_key_foreign_column(
-        ForeignKey(Pipeline).make(ondelete="CASCADE")
+    pipeline_id: sa_orm.Mapped[PipelineID] = orm.primary_key_foreign_column(
+        orm.ForeignKey(Pipeline).make(ondelete="CASCADE")
     )
-    source_id: sa_orm.Mapped[SourceID] = primary_key_foreign_column(
-        ForeignKey(Source).make(ondelete="CASCADE")
+    source_id: sa_orm.Mapped[SourceID] = orm.primary_key_foreign_column(
+        orm.ForeignKey(Source).make(ondelete="CASCADE")
     )
 
 
 class FeatureView(
-    SoftDeleteMixin, BelongsToOrgMixin, BelongsToRoomMixin, Base, kw_only=True
+    orm.SoftDeleteMixin,
+    orm.BelongsToOrgMixin,
+    BelongsToRoomMixin,
+    orm.Base,
+    kw_only=True,
 ):
     """A FeatureView is a logical collection of sources used by various spaces."""
 
     __tablename__ = "feature_view"
-    __table_args__ = (live_unique_constraint("name", "room_id"),)
+    __table_args__ = (orm.live_unique_constraint("name", "room_id"),)
 
-    id: sa_orm.Mapped[FeatureViewID | None] = primary_key_identity_column()
+    id: sa_orm.Mapped[FeatureViewID | None] = orm.primary_key_identity_column()
     name: sa_orm.Mapped[str] = sa_orm.mapped_column(sa.Text, default=None)
     description: sa_orm.Mapped[str] = sa_orm.mapped_column(sa.Text, default="")
 
@@ -240,22 +268,24 @@ class FeatureView(
     )
 
 
-class FeatureViewSource(BelongsToOrgMixin, BelongsToRoomMixin, Base, kw_only=True):
+class FeatureViewSource(
+    orm.BelongsToOrgMixin, BelongsToRoomMixin, orm.Base, kw_only=True
+):
     """A source inside of a feature view."""
 
     __tablename__ = "feature_view_source"
 
     table_op_graph: sa_orm.Mapped[table_pb2.TableComputeOp] = sa_orm.mapped_column()
     feature_view_id: sa_orm.Mapped[FeatureViewID] = sa_orm.mapped_column(
-        ForeignKey(FeatureView).make(ondelete="CASCADE"),
+        orm.ForeignKey(FeatureView).make(ondelete="CASCADE"),
         nullable=False,
     )
-    id: sa_orm.Mapped[FeatureViewSourceID | None] = primary_key_identity_column()
+    id: sa_orm.Mapped[FeatureViewSourceID | None] = orm.primary_key_identity_column()
     drop_disconnected: sa_orm.Mapped[bool] = sa_orm.mapped_column(default=False)
     # this should be legal but pyright complains that it makes Source depend
     # on itself
     source_id: sa_orm.Mapped[SourceID] = sa_orm.mapped_column(
-        ForeignKey(Source).make(ondelete="CASCADE"),
+        orm.ForeignKey(Source).make(ondelete="CASCADE"),
         nullable=False,
         default=None,
     )
@@ -264,18 +294,18 @@ class FeatureViewSource(BelongsToOrgMixin, BelongsToRoomMixin, Base, kw_only=Tru
     )
 
 
-class Space(BelongsToOrgMixin, BelongsToRoomMixin, Base, kw_only=True):
+class Space(orm.BelongsToOrgMixin, BelongsToRoomMixin, orm.Base, kw_only=True):
     """A space is a named evaluation of space parameters."""
 
     __tablename__ = "space"
     __table_args__ = (sa.UniqueConstraint("name", "room_id"),)
 
-    id: sa_orm.Mapped[SpaceID | None] = primary_key_identity_column()
+    id: sa_orm.Mapped[SpaceID | None] = orm.primary_key_identity_column()
     name: sa_orm.Mapped[str] = sa_orm.mapped_column(sa.Text, default=None)
     description: sa_orm.Mapped[str] = sa_orm.mapped_column(sa.Text, default="")
 
     feature_view_id: sa_orm.Mapped[FeatureViewID] = sa_orm.mapped_column(
-        ForeignKey(FeatureView).make(ondelete="CASCADE"),
+        orm.ForeignKey(FeatureView).make(ondelete="CASCADE"),
         nullable=False,
         default=None,
     )
@@ -294,13 +324,15 @@ class Space(BelongsToOrgMixin, BelongsToRoomMixin, Base, kw_only=True):
         return self.name
 
 
-class CompletionModel(SoftDeleteMixin, BelongsToOrgMixin, Base, kw_only=True):
+class CompletionModel(
+    orm.SoftDeleteMixin, orm.BelongsToOrgMixin, orm.Base, kw_only=True
+):
     """A customer's custom completion model definition."""
 
     __tablename__ = "completion_model"
-    __table_args__ = (live_unique_constraint("name", "org_id"),)
+    __table_args__ = (orm.live_unique_constraint("name", "org_id"),)
 
-    id: sa_orm.Mapped[CompletionModelID | None] = primary_key_identity_column()
+    id: sa_orm.Mapped[CompletionModelID | None] = orm.primary_key_identity_column()
     name: sa_orm.Mapped[str] = sa_orm.mapped_column(sa.Text, default=None)
     description: sa_orm.Mapped[str] = sa_orm.mapped_column(sa.Text, default=None)
     parameters: sa_orm.Mapped[completion_model_pb2.CompletionModelParameters | None] = (
@@ -324,43 +356,37 @@ class CompletionModel(SoftDeleteMixin, BelongsToOrgMixin, Base, kw_only=True):
 
 
 ID = (
-    AgentID
-    | AgentMessageID
+    orm.ID
     | CompletionModelID
     | FeatureViewID
     | FeatureViewSourceID
-    | MessageEntryID
-    | OrgID
     | PipelineID
     | ResourceID
     | RoomID
     | SourceID
     | SpaceID
-    | SpaceParametersID
-    | SpaceRunID
-    | UserMessageID
 )
 
+# These are part of the public interface, exposing them here so that
+# users don't have to import orm anytime they import eorm
+Base: TypeAlias = orm.Base
+Org: TypeAlias = orm.Org
+Session: TypeAlias = orm.Session
+
+InvalidORMIdentifierError: TypeAlias = orm.InvalidORMIdentifierError
+RequestedObjectsForNobodyError: TypeAlias = orm.RequestedObjectsForNobodyError
 
 __all__ = [
-    "AgentID",
-    "AgentMessageID",
     "Base",
-    "BaseID",
-    "BaseIDFromInt",
-    "BelongsToOrgMixin",
     "CompletionModel",
     "CompletionModelID",
     "DefaultObjects",
-    "DeletedObjectError",
     "FeatureView",
     "FeatureViewID",
     "FeatureViewSource",
     "FeatureViewSourceID",
     "ID",
     "InvalidORMIdentifierError",
-    "MessageEntryID",
-    "Org",
     "OrgID",
     "PipelineID",
     "PipelineInput",
@@ -375,12 +401,4 @@ __all__ = [
     "SourceID",
     "Space",
     "SpaceID",
-    "SpaceParametersID",
-    "SpaceRunID",
-    "UserMessageID",
-    "primary_key_foreign_column",
-    "primary_key_identity_column",
-    "primary_key_uuid_column",
-    "ProtoMessageDecorator",
-    "IntIDDecorator",
 ]