corvic-engine 0.3.0rc55__cp38-abi3-win_amd64.whl → 0.3.0rc57__cp38-abi3-win_amd64.whl

corvic/model/_proto_orm_convert.py

@@ -1,6 +1,6 @@
 import datetime
 from collections.abc import Callable, Sequence
-from typing import Any, Protocol, TypeVar, cast
+from typing import Any, Protocol, TypeVar
 
 import sqlalchemy as sa
 import sqlalchemy.orm as sa_orm
@@ -14,26 +14,30 @@ from corvic_generated.orm.v1 import feature_view_pb2
 
 UNCOMMITTED_ID_PREFIX = "__uncommitted_object-"
 
-_Proto = TypeVar(
-    "_Proto",
-    models_pb2.Resource,
-    models_pb2.Source,
-    models_pb2.FeatureView,
-    models_pb2.Space,
-    models_pb2.FeatureViewSource,
-    models_pb2.Agent,
-    models_pb2.Pipeline,
-    models_pb2.Room,
-    models_pb2.CompletionModel,
-)
 ID = TypeVar("ID", bound=orm.BaseID[Any])
 
 
+class _ModelProto(Protocol):
+    id: str
+    created_at: timestamp_pb2.Timestamp
+
+
+class _ModelBelongsToOrgProto(_ModelProto, Protocol):
+    org_id: str
+
+
+class _ModelBelongsToRoomProto(_ModelBelongsToOrgProto, Protocol):
+    room_id: str
+
+
+ProtoObj = TypeVar("ProtoObj", bound=_ModelProto)
+ProtoBelongsToOrgObj = TypeVar("ProtoBelongsToOrgObj", bound=_ModelBelongsToOrgProto)
+ProtoBelongsToRoomObj = TypeVar("ProtoBelongsToRoomObj", bound=_ModelBelongsToRoomProto)
+
+
 class _OrmModel(Protocol[ID]):
     id: sa_orm.Mapped[ID | None]
 
-    org_id: sa_orm.Mapped[orm.OrgID | None]
-
     @sa.ext.hybrid.hybrid_property
     def created_at(self) -> datetime.datetime | None: ...
 
@@ -42,7 +46,17 @@ class _OrmModel(Protocol[ID]):
     def _created_at_expression(cls): ...
 
 
+class _OrmBelongsToOrgModel(_OrmModel[ID], Protocol):
+    org_id: sa_orm.Mapped[orm.OrgID | None]
+
+
+class _OrmBelongsToRoomModel(_OrmBelongsToOrgModel[ID], Protocol):
+    room_id: sa_orm.Mapped[orm.RoomID | None]
+
+
 OrmObj = TypeVar("OrmObj", bound=_OrmModel[Any])
+OrmBelongsToOrgObj = TypeVar("OrmBelongsToOrgObj", bound=_OrmBelongsToOrgModel[Any])
+OrmBelongsToRoomObj = TypeVar("OrmBelongsToRoomObj", bound=_OrmBelongsToRoomModel[Any])
 
 
 def _translate_orm_id(
@@ -58,39 +72,6 @@ def _translate_orm_id(
             return Ok(parsed_obj_id)
 
 
-def _translate_orm_ids(
-    proto_obj: _Proto, obj_id_class: type[ID]
-) -> Ok[tuple[ID | None, orm.RoomID | None]] | orm.InvalidORMIdentifierError:
-    match _translate_orm_id(proto_obj.id, obj_id_class):
-        case orm.InvalidORMIdentifierError() as err:
-            return err
-        case Ok(obj_id):
-            pass
-
-    match proto_obj:
-        case (
-            models_pb2.Resource()
-            | models_pb2.Source()
-            | models_pb2.FeatureView()
-            | models_pb2.Space()
-            | models_pb2.Agent()
-            | models_pb2.Pipeline()
-            | models_pb2.FeatureViewSource()
-        ):
-            room_id = orm.RoomID(proto_obj.room_id)
-            match room_id.to_db():
-                case orm.InvalidORMIdentifierError() as err:
-                    return err
-                case Ok():
-                    pass
-        case models_pb2.CompletionModel():
-            room_id = None
-        case models_pb2.Room():
-            room_id = cast(orm.RoomID, obj_id)
-
-    return Ok((obj_id, room_id))
-
-
 def timestamp_orm_to_proto(
     timestamp_orm: datetime.datetime | None,
 ) -> timestamp_pb2.Timestamp | None:
@@ -247,43 +228,51 @@ def completion_model_orm_to_proto(
     )
 
 
-def _add_orm_to_session(
-    orm_obj: OrmObj, org_id: str, session: sa_orm.Session
-) -> OrmObj:
-    if org_id:
-        orm_obj.org_id = orm.OrgID(org_id)
+def add_orm_org_mixin_to_session(
+    orm_obj: OrmBelongsToOrgObj,
+    proto_obj: _ModelBelongsToOrgProto,
+    id_class: type[ID],
+    session: sa_orm.Session,
+) -> Ok[OrmBelongsToOrgObj] | orm.InvalidORMIdentifierError:
+    match _translate_orm_id(proto_obj.id, id_class):
+        case Ok(orm_id):
+            orm_obj.id = orm_id
+        case orm.InvalidORMIdentifierError() as err:
+            return err
+    if proto_obj.org_id:
+        org_id = orm.OrgID(proto_obj.org_id)
+        match org_id.to_db():
+            case Ok():
+                orm_obj.org_id = org_id
+            case orm.InvalidORMIdentifierError() as err:
+                return err
+        orm_obj.org_id = org_id
     if not orm_obj.id:
         session.add(orm_obj)
     else:
         orm_obj = session.merge(orm_obj)
-    return orm_obj
+    return Ok(orm_obj)
 
 
-def resource_proto_to_orm(
-    proto_obj: models_pb2.Resource, session: sa_orm.Session
-) -> Ok[orm.Resource] | InvalidArgumentError:
-    match _translate_orm_ids(proto_obj, orm.ResourceID):
+def add_orm_room_mixin_to_session(
+    orm_obj: OrmBelongsToRoomObj,
+    proto_obj: _ModelBelongsToRoomProto,
+    id_class: type[ID],
+    session: sa_orm.Session,
+) -> Ok[OrmBelongsToRoomObj] | orm.InvalidORMIdentifierError:
+    room_id = orm.RoomID(proto_obj.room_id)
+    match room_id.to_db():
+        case Ok():
+            pass
         case orm.InvalidORMIdentifierError() as err:
             return err
-        case Ok((obj_id, room_id)):
-            pass
-    if not room_id:
-        return InvalidArgumentError("room id required to commit")
+    orm_obj.room_id = orm.RoomID(proto_obj.room_id)
+    return add_orm_org_mixin_to_session(orm_obj, proto_obj, id_class, session)
 
-    orm_obj = orm.Resource(
-        id=obj_id,
-        name=proto_obj.name,
-        description=proto_obj.description,
-        mime_type=proto_obj.mime_type,
-        md5=proto_obj.md5,
-        url=proto_obj.url,
-        size=proto_obj.size,
-        original_path=proto_obj.original_path,
-        latest_event=proto_obj.recent_events[-1] if proto_obj.recent_events else None,
-        room_id=room_id,
-    )
-    _add_orm_to_session(orm_obj, proto_obj.org_id, session)
 
+def _resource_pipeline_to_orm(
+    proto_obj: models_pb2.Resource, orm_obj: orm.Resource, session: sa_orm.Session
+) -> Ok[None] | InvalidArgumentError:
     if proto_obj.pipeline_id:
         match _translate_orm_id(proto_obj.pipeline_id, orm.PipelineID):
             case orm.InvalidORMIdentifierError() as err:
@@ -299,18 +288,41 @@ def resource_proto_to_orm(
             resource_id=orm_obj.id,
             name=proto_obj.pipeline_input_name,
             pipeline_id=pipeline_id,
-            room_id=room_id,
+            room_id=orm_obj.room_id,
         )
         if orm_obj.org_id:
             pipeline_input.org_id = orm_obj.org_id
         orm_obj.pipeline_ref = session.merge(pipeline_input)
+    return Ok(None)
 
-    return Ok(orm_obj)
+
+def resource_proto_to_orm(
+    proto_obj: models_pb2.Resource, session: sa_orm.Session
+) -> Ok[orm.Resource] | InvalidArgumentError:
+    orm_obj = orm.Resource(
+        name=proto_obj.name,
+        description=proto_obj.description,
+        mime_type=proto_obj.mime_type,
+        md5=proto_obj.md5,
+        url=proto_obj.url,
+        size=proto_obj.size,
+        original_path=proto_obj.original_path,
+        latest_event=proto_obj.recent_events[-1] if proto_obj.recent_events else None,
+    )
+    add_orm_room_mixin_to_session(orm_obj, proto_obj, orm.ResourceID, session)
+
+    match _resource_pipeline_to_orm(proto_obj, orm_obj, session):
+        case Ok(None):
+            return Ok(orm_obj)
+        case InvalidArgumentError() as err:
+            return err
 
 
 def _ensure_id(
-    proto_obj: _Proto,
-    proto_to_orm: Callable[[_Proto, sa_orm.Session], Ok[OrmObj] | InvalidArgumentError],
+    proto_obj: ProtoObj,
+    proto_to_orm: Callable[
+        [ProtoObj, sa_orm.Session], Ok[OrmObj] | InvalidArgumentError
+    ],
     id_type: type[ID],
     session: sa_orm.Session,
 ) -> Ok[ID] | orm.InvalidORMIdentifierError | InvalidArgumentError:
@@ -333,23 +345,19 @@ def _ensure_id(
 def pipeline_proto_to_orm(  # noqa: C901
     proto_obj: models_pb2.Pipeline, session: sa_orm.Session
 ) -> Ok[orm.Pipeline] | orm.InvalidORMIdentifierError | InvalidArgumentError:
-    match _translate_orm_ids(proto_obj, orm.PipelineID):
-        case orm.InvalidORMIdentifierError() as err:
-            return err
-        case Ok((obj_id, room_id)):
-            pass
-    if not room_id:
-        return InvalidArgumentError("room id required to commit")
-
     inputs = list[orm.PipelineInput]()
     orm_obj = orm.Pipeline(
-        id=obj_id,
         name=proto_obj.name,
-        room_id=room_id,
         transformation=proto_obj.pipeline_transformation,
         description=proto_obj.description,
     )
-    orm_obj = _add_orm_to_session(orm_obj, proto_obj.org_id, session)
+    if proto_obj.org_id:
+        orm_obj.org_id = orm.OrgID(proto_obj.org_id)
+    match add_orm_room_mixin_to_session(orm_obj, proto_obj, orm.PipelineID, session):
+        case Ok(orm_obj):
+            pass
+        case orm.InvalidORMIdentifierError() as err:
+            return err
     session.flush()
 
     if not orm_obj.id:
@@ -364,7 +372,10 @@ def pipeline_proto_to_orm(  # noqa: C901
                 pass
         outputs.append(
             orm.PipelineOutput(
-                source_id=source_id, name=name, pipeline_id=orm_obj.id, room_id=room_id
+                source_id=source_id,
+                name=name,
+                pipeline_id=orm_obj.id,
+                room_id=orm_obj.room_id,
             )
         )
 
@@ -380,21 +391,11 @@ def pipeline_proto_to_orm(  # noqa: C901
 def source_proto_to_orm(
     proto_obj: models_pb2.Source, session: sa_orm.Session
 ) -> Ok[orm.Source] | orm.InvalidORMIdentifierError | InvalidArgumentError:
-    match _translate_orm_ids(proto_obj, orm.SourceID):
-        case orm.InvalidORMIdentifierError() as err:
-            return err
-        case Ok((obj_id, room_id)):
-            pass
-    if not room_id:
-        return InvalidArgumentError("room id required to commit")
-
     orm_obj = orm.Source(
-        id=obj_id,
         name=proto_obj.name,
         table_op_graph=proto_obj.table_op_graph,
-        room_id=room_id,
     )
-    return Ok(_add_orm_to_session(orm_obj, proto_obj.org_id, session))
+    return add_orm_room_mixin_to_session(orm_obj, proto_obj, orm.SourceID, session)
 
 
 def _update_agent_associations(
@@ -431,21 +432,15 @@ def _update_agent_associations(
 def agent_proto_to_orm(
     proto_obj: models_pb2.Agent, session: sa_orm.Session
 ) -> Ok[orm.Agent] | orm.InvalidORMIdentifierError | InvalidArgumentError:
-    match _translate_orm_ids(proto_obj, orm.AgentID):
-        case orm.InvalidORMIdentifierError() as err:
-            return err
-        case Ok((obj_id, room_id)):
-            pass
-    if not room_id:
-        return InvalidArgumentError("room id required to commit")
-
     orm_obj = orm.Agent(
-        id=obj_id,
         name=proto_obj.name,
         parameters=proto_obj.agent_parameters,
-        room_id=room_id,
     )
-    agent = _add_orm_to_session(orm_obj, proto_obj.org_id, session)
+    match add_orm_room_mixin_to_session(orm_obj, proto_obj, orm.AgentID, session):
+        case Ok(agent):
+            pass
+        case orm.InvalidORMIdentifierError() as err:
+            return err
     session.flush()
 
     if not agent.id:
@@ -462,14 +457,6 @@ def agent_proto_to_orm(
 def space_proto_to_orm(
     proto_obj: models_pb2.Space, session: sa_orm.Session
 ) -> Ok[orm.Space] | orm.InvalidORMIdentifierError | InvalidArgumentError:
-    match _translate_orm_ids(proto_obj, orm.SpaceID):
-        case orm.InvalidORMIdentifierError() as err:
-            return err
-        case Ok((obj_id, room_id)):
-            pass
-    if not room_id:
-        return InvalidArgumentError("room id required to commit")
-
     match _ensure_id(
         proto_obj.feature_view, feature_view_proto_to_orm, orm.FeatureViewID, session
     ):
@@ -482,38 +469,29 @@ def space_proto_to_orm(
         raise InternalError("internal assertion did not hold")
 
     orm_obj = orm.Space(
-        id=obj_id,
         name=proto_obj.name,
         description=proto_obj.description,
-        room_id=room_id,
         feature_view_id=feature_view_id,
         parameters=proto_obj.space_parameters,
         auto_sync=proto_obj.auto_sync,
     )
-    return Ok(_add_orm_to_session(orm_obj, proto_obj.org_id, session))
+    return add_orm_room_mixin_to_session(orm_obj, proto_obj, orm.SpaceID, session)
 
 
-def feature_view_proto_to_orm(  # noqa: C901
+def feature_view_proto_to_orm(
     proto_obj: models_pb2.FeatureView, session: sa_orm.Session
 ) -> Ok[orm.FeatureView] | orm.InvalidORMIdentifierError | InvalidArgumentError:
-    match _translate_orm_ids(proto_obj, orm.FeatureViewID):
-        case orm.InvalidORMIdentifierError() as err:
-            return err
-        case Ok((obj_id, room_id)):
-            pass
-    if not room_id:
-        return InvalidArgumentError("room id required to commit")
-
     orm_obj = orm.FeatureView(
-        id=obj_id,
         name=proto_obj.name,
         description=proto_obj.description,
-        room_id=room_id,
     )
     if proto_obj.org_id:
         orm_obj.org_id = orm.OrgID(proto_obj.org_id)
-
-    orm_obj = _add_orm_to_session(orm_obj, proto_obj.org_id, session)
+    match add_orm_room_mixin_to_session(orm_obj, proto_obj, orm.FeatureViewID, session):
+        case Ok(orm_obj):
+            pass
+        case orm.InvalidORMIdentifierError() as err:
+            return err
     session.flush()
 
     if not orm_obj.id:
@@ -573,57 +551,38 @@ def _feature_view_source_proto_to_orm(
             return err
         case Ok(source_id):
             pass
-    match _translate_orm_id(proto_obj.id, orm.FeatureViewSourceID):
-        case orm.InvalidORMIdentifierError() as err:
-            return err
-        case Ok(obj_id):
-            pass
 
     orm_obj = orm.FeatureViewSource(
-        room_id=orm.RoomID(proto_obj.room_id),
-        id=obj_id,
         table_op_graph=proto_obj.table_op_graph,
         drop_disconnected=proto_obj.drop_disconnected,
         source_id=source_id,
         feature_view_id=feature_view_id,
     )
-    return Ok(_add_orm_to_session(orm_obj, proto_obj.org_id, session))
+    return add_orm_room_mixin_to_session(
+        orm_obj, proto_obj, orm.FeatureViewSourceID, session
+    )
 
 
 def room_proto_to_orm(
     proto_obj: models_pb2.Room, session: sa_orm.Session
 ) -> Ok[orm.Room] | orm.InvalidORMIdentifierError | InvalidArgumentError:
-    match _translate_orm_ids(proto_obj, orm.RoomID):
-        case orm.InvalidORMIdentifierError() as err:
-            return err
-        case Ok((obj_id, _)):
-            pass
-
-    orm_obj = orm.Room(
-        id=obj_id,
-        name=proto_obj.name,
-    )
-    return Ok(_add_orm_to_session(orm_obj, proto_obj.org_id, session))
+    orm_obj = orm.Room(name=proto_obj.name)
+    return add_orm_org_mixin_to_session(orm_obj, proto_obj, orm.RoomID, session)
 
 
 def completion_model_proto_to_orm(
     proto_obj: models_pb2.CompletionModel, session: sa_orm.Session
 ) -> Ok[orm.CompletionModel] | InvalidArgumentError:
-    match _translate_orm_ids(proto_obj, orm.CompletionModelID):
-        case orm.InvalidORMIdentifierError() as err:
-            return err
-        case Ok((obj_id, _)):
-            pass
-
     orm_obj = orm.CompletionModel(
-        id=obj_id,
         name=proto_obj.name,
         description=proto_obj.description,
         parameters=proto_obj.parameters,
         secret_api_key=proto_obj.secret_api_key,
         last_validation_time=proto_obj.last_validation_time.ToDatetime(),
     )
-    return Ok(_add_orm_to_session(orm_obj, proto_obj.org_id, session))
+    return add_orm_org_mixin_to_session(
+        orm_obj, proto_obj, orm.CompletionModelID, session
+    )
 
 
 def source_delete_orms(

corvic/model/_resource.py

@@ -15,7 +15,7 @@ from sqlalchemy.orm.interfaces import LoaderOption
 from typing_extensions import Self
 
 from corvic import orm, system
-from corvic.model._base_model import BaseModel
+from corvic.model._base_model import BelongsToRoomModel
 from corvic.model._defaults import Defaults
 from corvic.model._proto_orm_convert import (
     resource_delete_orms,
@@ -32,7 +32,7 @@ RoomID: TypeAlias = orm.RoomID
 PipelineID: TypeAlias = orm.PipelineID
 
 
-class Resource(BaseModel[ResourceID, models_pb2.Resource, orm.Resource]):
+class Resource(BelongsToRoomModel[ResourceID, models_pb2.Resource, orm.Resource]):
     """Resources represent import data."""
 
     @classmethod
@@ -67,10 +67,6 @@ class Resource(BaseModel[ResourceID, models_pb2.Resource, orm.Resource]):
     def name(self) -> str:
         return self.proto_self.name
 
-    @property
-    def room_id(self) -> RoomID:
-        return RoomID(self.proto_self.room_id)
-
     @property
     def pipeline_id(self) -> PipelineID | None:
         return PipelineID(self.proto_self.pipeline_id) or None

corvic/model/_room.py

@@ -8,7 +8,7 @@ from typing import TypeAlias
 import structlog
 
 from corvic import orm, system
-from corvic.model._base_model import BaseModel
+from corvic.model._base_model import BelongsToOrgModel
 from corvic.model._defaults import Defaults
 from corvic.model._proto_orm_convert import (
     room_delete_orms,
@@ -25,7 +25,7 @@ RoomID: TypeAlias = orm.RoomID
 FeatureViewID: TypeAlias = orm.FeatureViewID
 
 
-class Room(BaseModel[RoomID, models_pb2.Room, orm.Room]):
+class Room(BelongsToOrgModel[RoomID, models_pb2.Room, orm.Room]):
     """Rooms contain conversations and tables."""
 
     @classmethod

corvic/model/_source.py

@@ -15,7 +15,7 @@ from sqlalchemy.orm.interfaces import LoaderOption
 from typing_extensions import Self
 
 from corvic import op_graph, orm, system
-from corvic.model._base_model import BaseModel
+from corvic.model._base_model import BelongsToRoomModel
 from corvic.model._defaults import Defaults
 from corvic.model._proto_orm_convert import (
     source_delete_orms,
@@ -46,7 +46,7 @@ def foreign_key(
             )
 
 
-class Source(BaseModel[SourceID, models_pb2.Source, orm.Source]):
+class Source(BelongsToRoomModel[SourceID, models_pb2.Source, orm.Source]):
     """Sources describe how resources should be treated.
 
     Example:
@@ -269,14 +269,18 @@ class Source(BaseModel[SourceID, models_pb2.Source, orm.Source]):
             self.client, op_graph.op.from_proto(self.proto_self.table_op_graph)
         )
 
+    @functools.cached_property
+    def prop_table(self):
+        if self.proto_self.prop_table_op_graph is None:
+            return None
+        return Table.from_ops(
+            self.client, op_graph.op.from_proto(self.proto_self.prop_table_op_graph)
+        )
+
     @property
     def name(self) -> str:
         return self.proto_self.name
 
-    @property
-    def room_id(self) -> RoomID:
-        return RoomID(self.proto_self.room_id)
-
     @property
     def pipeline_id(self) -> PipelineID | None:
         return PipelineID(self.proto_self.pipeline_id) or None

corvic/model/_space.py

@@ -14,7 +14,7 @@ from sqlalchemy import orm as sa_orm
 from typing_extensions import Self
 
 from corvic import op_graph, orm, system
-from corvic.model._base_model import BaseModel
+from corvic.model._base_model import BelongsToRoomModel
 from corvic.model._defaults import Defaults
 from corvic.model._feature_view import FeatureView, FeatureViewEdgeTableMetadata
 from corvic.model._proto_orm_convert import (
@@ -61,7 +61,7 @@ name_to_proto_image_model = {
 }
 
 
-class Space(BaseModel[SpaceID, models_pb2.Space, orm.Space]):
+class Space(BelongsToRoomModel[SpaceID, models_pb2.Space, orm.Space]):
     """Spaces apply embedding methods to FeatureViews.
 
     Example:
@@ -105,10 +105,6 @@ class Space(BaseModel[SpaceID, models_pb2.Space, orm.Space]):
     def name(self):
         return self.proto_self.name
 
-    @property
-    def room_id(self):
-        return RoomID(self.proto_self.room_id)
-
     @property
     def description(self):
         return self.proto_self.description
@@ -840,13 +836,20 @@ class TabularSpace(Space):
                     pa.string(),
                 )
 
-            op = (
-                op.concat_list(
+            if len(embedding_column_tmp_names):
+                op = op.concat_list(
                     column_names=embedding_column_tmp_names,
                     concat_list_column_name=embedding_column_tmp_name,
+                ).and_then(reduce_dimension)
+            else:
+                op = op.add_literal_column(
+                    column_name=embedding_column_tmp_name,
+                    literal=[],
+                    dtype=pa.large_list(pa.float32()),
+                    ftype=op_graph.feature_type.embedding(),
                 )
-                .and_then(reduce_dimension)
-                .and_then(select_columns)
+            op = (
+                op.and_then(select_columns)
                 .and_then(update_feature_types)
                 .and_then(rename_columns)
                 .and_then(add_literal_column)

corvic/op_graph/_schema.py

@@ -2,8 +2,6 @@
 
 from __future__ import annotations
 
-from collections.abc import Sequence
-
 import polars as pl
 import pyarrow as pa
 
@@ -14,7 +12,7 @@ from corvic.op_graph.feature_types import FeatureType
 
 def _infer_possible_feature_types(
     data_type: pa.DataType,
-) -> Sequence[FeatureType]:
+) -> list[FeatureType]:
     if pa.types.is_integer(data_type):
         # TODO(thunt): Add multi-categorical in future versions
         return [
@@ -168,6 +166,6 @@ class Field:
     def rename(self, new_name: str) -> Field:
         return Field(new_name, self.dtype, self.ftype)
 
-    def possible_feature_types(self) -> Sequence[FeatureType]:
+    def possible_feature_types(self) -> list[FeatureType]:
         """Infer possible feature types given the data type."""
         return _infer_possible_feature_types(self.dtype)