corvic-engine 0.3.0rc42__cp38-abi3-win_amd64.whl → 0.3.0rc44__cp38-abi3-win_amd64.whl

corvic/embedding_metric/embeddings.py

@@ -84,8 +84,21 @@ def stable_rank(
         case Ok():
             pass
 
-    numerator = norm(embeddings, order="fro")
-    denominator = norm(embeddings, order=2) ** 2
+    try:
+        numerator = norm(embeddings, order="fro")
+        denominator = norm(embeddings, order=2) ** 2
+    except np.linalg.LinAlgError as err:
+        return InvalidArgumentError.from_(err)
+
+    if not np.isfinite(numerator):
+        return InvalidArgumentError("embeddings norm_fro is not finite")
+
+    if not np.isfinite(denominator):
+        return InvalidArgumentError("embeddings norm_2 is not finite")
+
+    if not denominator:
+        return InvalidArgumentError("embeddings norm_2 is zero")
+
     metric = numerator / denominator
     if normalize:
         metric = 1 - 1 / (1 + metric)
@@ -162,7 +175,10 @@ def ne_sum(
             pass
 
     covariance = np.cov(embeddings.T)
-    eigenvalues = linalg.eigvals(covariance)
+    try:
+        eigenvalues = linalg.eigvals(covariance)
+    except np.linalg.LinAlgError as err:
+        return InvalidArgumentError.from_(err)
 
     # Discard imaginary part
     eigenvalues = eigenvalues.real
@@ -172,7 +188,7 @@ def ne_sum(
     eigenvalues = eigenvalues[sorted_indices]
 
     if eigenvalues[0] == 0:
-        return Ok(0)
+        return Ok(0.0)
 
     ne_sum_value = float(np.sum(eigenvalues) / eigenvalues[0])
     if normalize:
@@ -227,7 +243,11 @@ def condition_number(
         case Ok():
             pass
 
-    metric = float(np.linalg.cond(embeddings, p=p))
+    try:
+        metric = float(np.linalg.cond(embeddings, p=p))
+    except np.linalg.LinAlgError as err:
+        return InvalidArgumentError.from_(err)
+
     if normalize:
         metric = 1 - 1 / (1 + metric)
     return Ok(metric)
@@ -246,7 +266,11 @@ def rcondition_number(
         case Ok():
             pass
 
-    metric = float(1 / np.linalg.cond(embeddings, p=p))
+    try:
+        metric = float(1 / np.linalg.cond(embeddings, p=p))
+    except np.linalg.LinAlgError as err:
+        return InvalidArgumentError.from_(err)
+
     if normalize:
         metric = 1 - 1 / (1 + metric)
     return Ok(metric)

corvic/model/_completion_model.py

@@ -5,7 +5,7 @@ from __future__ import annotations
 import copy
 import datetime
 from collections.abc import Iterable, Sequence
-from typing import TypeAlias
+from typing import Literal, TypeAlias
 
 from sqlalchemy import orm as sa_orm
 
@@ -19,6 +19,7 @@ from corvic.model._proto_orm_convert import (
 )
 from corvic.result import InvalidArgumentError, NotFoundError, Ok
 from corvic_generated.model.v1alpha import models_pb2
+from corvic_generated.orm.v1 import completion_model_pb2
 
 CompletionModelID: TypeAlias = orm.CompletionModelID
 OrgID: TypeAlias = orm.OrgID
@@ -62,12 +63,46 @@ class CompletionModel(
         return OrgID(self.proto_self.org_id)
 
     @property
-    def model_name(self) -> str:
-        return self.proto_self.model_name
+    def provider(self) -> Literal["openai-generic", "azure-openai"] | None:
+        match self.proto_self.parameters.WhichOneof("params"):
+            case "azure_openai_parameters":
+                return "azure-openai"
+            case "generic_openai_parameters":
+                return "openai-generic"
+            case _:
+                return None
 
     @property
-    def endpoint(self) -> str:
-        return self.proto_self.endpoint
+    def parameters(
+        self,
+    ) -> (
+        completion_model_pb2.AzureOpenAIParameters
+        | completion_model_pb2.GenericOpenAIParameters
+        | None
+    ):
+        match self.provider:
+            case "azure-openai":
+                return self.azure_openai_parameters
+            case "openai-generic":
+                return self.generic_openai_parameters
+            case None:
+                return None
+
+    @property
+    def azure_openai_parameters(
+        self,
+    ) -> completion_model_pb2.AzureOpenAIParameters | None:
+        if self.proto_self.parameters.HasField("azure_openai_parameters"):
+            return self.proto_self.parameters.azure_openai_parameters
+        return None
+
+    @property
+    def generic_openai_parameters(
+        self,
+    ) -> completion_model_pb2.GenericOpenAIParameters | None:
+        if self.proto_self.parameters.HasField("generic_openai_parameters"):
+            return self.proto_self.parameters.generic_openai_parameters
+        return None
 
     @property
     def secret_api_key(self) -> str:
@@ -83,8 +118,7 @@ class CompletionModel(
         *,
         name: str,
         description: str,
-        model_name: str,
-        endpoint: str,
+        parameters: completion_model_pb2.CompletionModelParameters,
         secret_api_key: str,
         client: system.Client | None = None,
     ):
@@ -94,8 +128,7 @@ class CompletionModel(
             models_pb2.CompletionModel(
                 name=name,
                 description=description,
-                model_name=model_name,
-                endpoint=endpoint,
+                parameters=parameters,
                 secret_api_key=secret_api_key,
             ),
         )
@@ -153,14 +186,11 @@ class CompletionModel(
         proto_self.description = description
         return CompletionModel(self.client, proto_self)
 
-    def with_model_name(self, model_name: str) -> CompletionModel:
-        proto_self = copy.deepcopy(self.proto_self)
-        proto_self.model_name = model_name
-        return CompletionModel(self.client, proto_self)
-
-    def with_endpoint(self, endpoint: str) -> CompletionModel:
+    def with_parameters(
+        self, parameters: completion_model_pb2.CompletionModelParameters
+    ) -> CompletionModel:
         proto_self = copy.deepcopy(self.proto_self)
-        proto_self.endpoint = endpoint
+        proto_self.parameters.CopyFrom(parameters)
         return CompletionModel(self.client, proto_self)
 
     def with_secret_api_key(self, secret_api_key: str) -> CompletionModel:

corvic/model/_feature_view.py

@@ -1075,6 +1075,7 @@ class FeatureView(BaseModel[FeatureViewID, models_pb2.FeatureView, orm.FeatureVi
             )
 
         proto_feature_view_source = models_pb2.FeatureViewSource(
+            room_id=str(source.room_id),
             table_op_graph=new_table.op_graph.to_proto(),
             drop_disconnected=drop_disconnected,
             source=source.proto_self,

corvic/model/_proto_orm_convert.py

@@ -66,6 +66,7 @@ def _translate_orm_ids(
             | 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():
@@ -73,7 +74,7 @@ def _translate_orm_ids(
                     return err
                 case Ok():
                     pass
-        case models_pb2.FeatureViewSource() | models_pb2.CompletionModel():
+        case models_pb2.CompletionModel():
             room_id = None
         case models_pb2.Room():
             room_id = cast(orm.RoomID, obj_id)
@@ -146,6 +147,7 @@ def feature_view_source_orm_to_proto(
 ) -> models_pb2.FeatureViewSource:
     return models_pb2.FeatureViewSource(
         id=str(feature_view_source_orm.id),
+        room_id=str(feature_view_source_orm.room_id),
         source=source_orm_to_proto(feature_view_source_orm.source),
         table_op_graph=feature_view_source_orm.table_op_graph,
         drop_disconnected=feature_view_source_orm.drop_disconnected,
@@ -225,8 +227,7 @@ def completion_model_orm_to_proto(
         name=completion_model_orm.name,
         description=completion_model_orm.description,
         org_id=str(completion_model_orm.org_id),
-        model_name=completion_model_orm.model_name,
-        endpoint=completion_model_orm.endpoint,
+        parameters=completion_model_orm.parameters,
         secret_api_key=completion_model_orm.secret_api_key,
         created_at=timestamp_orm_to_proto(completion_model_orm.created_at),
     )
@@ -263,7 +264,9 @@ def resource_proto_to_orm(
         latest_event=proto_obj.recent_events[-1] if proto_obj.recent_events else None,
         room_id=ids.room_id,
         source_associations=[
-            orm.SourceResourceAssociation(source_id=src_id, resource_id=ids.obj_id)
+            orm.SourceResourceAssociation(
+                room_id=ids.room_id, source_id=src_id, resource_id=ids.obj_id
+            )
             for src_id in source_ids
         ],
     )
@@ -317,7 +320,10 @@ def pipeline_proto_to_orm(  # noqa: C901
                     session.add(resource_orm)
                 session.merge(
                     orm.PipelineInput(
-                        pipeline=orm_obj, resource=resource_orm, name=name
+                        room_id=resource_orm.room_id,
+                        pipeline=orm_obj,
+                        resource=resource_orm,
+                        name=name,
                     )
                 )
 
@@ -333,7 +339,12 @@ def pipeline_proto_to_orm(  # noqa: C901
                 else:
                     session.add(source_orm)
                 session.merge(
-                    orm.PipelineOutput(pipeline=orm_obj, source=source_orm, name=name)
+                    orm.PipelineOutput(
+                        room_id=source_orm.room_id,
+                        pipeline=orm_obj,
+                        source=source_orm,
+                        name=name,
+                    )
                 )
     if proto_obj.org_id:
         org_id = orm.OrgID(proto_obj.org_id)
@@ -356,7 +367,9 @@ def source_proto_to_orm(
     resource_id = orm.ResourceID(proto_obj.resource_id)
     if resource_id:
         associations = [
-            orm.SourceResourceAssociation(source_id=ids.obj_id, resource_id=resource_id)
+            orm.SourceResourceAssociation(
+                room_id=ids.room_id, source_id=ids.obj_id, resource_id=resource_id
+            )
         ]
     else:
         associations = list[orm.SourceResourceAssociation]()
@@ -499,6 +512,7 @@ def feature_view_source_proto_to_orm(
     else:
         session.add(source)
     orm_obj = orm.FeatureViewSource(
+        room_id=source.room_id,
         table_op_graph=proto_obj.table_op_graph,
         drop_disconnected=proto_obj.drop_disconnected,
         source=source,
@@ -543,8 +557,7 @@ def completion_model_proto_to_orm(
         id=ids.obj_id,
         name=proto_obj.name,
         description=proto_obj.description,
-        model_name=proto_obj.model_name,
-        endpoint=proto_obj.endpoint,
+        parameters=proto_obj.parameters,
         secret_api_key=proto_obj.secret_api_key,
     )
 

corvic/model/_space.py

@@ -782,50 +782,67 @@ class TabularSpace(Space):
 
             embedding_column_tmp_name = f"__embed-{uuid.uuid4()}"
 
-            op = (
-                op.concat_list(
-                    column_names=embedding_column_tmp_names,
-                    concat_list_column_name=embedding_column_tmp_name,
-                )
-                .and_then(
-                    lambda t,
-                    embedding_column_name=embedding_column_tmp_name: op_graph.op.coordinates_from_embedding(  # noqa: E501
-                        table=t,
-                        embedding_column_name=embedding_column_name,
-                        output_dims=parameters.ndim,
-                    )
+            # Avoid 0 padding for spaces with small numbers of columns
+            target_list_length = min(parameters.ndim, len(embedding_column_tmp_names))
+
+            def reduce_dimension(
+                op: op_graph.Op,
+                embedding_column_tmp_name=embedding_column_tmp_name,
+                target_list_length=target_list_length,
+            ):
+                return op.truncate_list(
+                    list_column_name=embedding_column_tmp_name,
+                    target_list_length=target_list_length,
+                    padding_value=0,
                 )
-                .and_then(
-                    lambda t,
-                    pk_field=pk_field,
-                    embedding_column_tmp_name=embedding_column_tmp_name,
-                    output_source=output_source: t.select_columns(
-                        [pk_field.name, embedding_column_tmp_name]
-                    )
+
+            def select_columns(
+                op: op_graph.Op,
+                pk_field=pk_field,
+                embedding_column_tmp_name=embedding_column_tmp_name,
+            ):
+                return op.select_columns([pk_field.name, embedding_column_tmp_name])
+
+            def update_feature_types(
+                op: op_graph.Op,
+                embedding_column_tmp_name=embedding_column_tmp_name,
+            ):
+                return op.update_feature_types(
+                    {embedding_column_tmp_name: op_graph.feature_type.embedding()}
                 )
-                .and_then(
-                    lambda t,
-                    embedding_column_tmp_name=embedding_column_tmp_name: t.update_feature_types(  # noqa: E501
-                        {embedding_column_tmp_name: op_graph.feature_type.embedding()}
-                    )
+
+            def rename_columns(
+                op: op_graph.Op,
+                pk_field=pk_field,
+                embedding_column_tmp_name=embedding_column_tmp_name,
+            ):
+                return op.rename_columns(
+                    {
+                        pk_field.name: "entity_id",
+                        embedding_column_tmp_name: "embedding",
+                    }
                 )
-                .and_then(
-                    lambda t,
-                    pk_field=pk_field,
-                    embedding_column_tmp_name=embedding_column_tmp_name: t.rename_columns(  # noqa: E501
-                        {
-                            pk_field.name: "entity_id",
-                            embedding_column_tmp_name: "embedding",
-                        }
-                    )
+
+            def add_literal_column(
+                op: op_graph.Op,
+                output_source=output_source,
+            ):
+                return op.add_literal_column(
+                    "source_id",
+                    str(output_source.id),
+                    pa.string(),
                 )
-                .and_then(
-                    lambda t, output_source=output_source: t.add_literal_column(
-                        "source_id",
-                        str(output_source.id),
-                        pa.string(),
-                    )
+
+            op = (
+                op.concat_list(
+                    column_names=embedding_column_tmp_names,
+                    concat_list_column_name=embedding_column_tmp_name,
                 )
+                .and_then(reduce_dimension)
+                .and_then(select_columns)
+                .and_then(update_feature_types)
+                .and_then(rename_columns)
+                .and_then(add_literal_column)
             )
 
             match op:

corvic/orm/__init__.py

@@ -48,6 +48,7 @@ from corvic.orm.mixins import (
 from corvic_generated.orm.v1 import (
     agent_pb2,
     common_pb2,
+    completion_model_pb2,
     feature_view_pb2,
     pipeline_pb2,
     space_pb2,
@@ -104,6 +105,13 @@ class Room(BelongsToOrgMixin, SoftDeleteMixin, Base):
         return self.name
 
 
+class BelongsToRoomMixin(sa_orm.MappedAsDataclass):
+    room_id: sa_orm.Mapped[RoomID | None] = sa_orm.mapped_column(
+        ForeignKey(Room).make(ondelete="CASCADE"),
+        nullable=True,
+    )
+
+
 class DefaultObjects(Base):
     """Holds the identifiers for default objects."""
 
@@ -117,7 +125,7 @@ class DefaultObjects(Base):
     version: sa_orm.Mapped[int | None] = primary_key_identity_column(type_=INT_PK_TYPE)
 
 
-class Resource(BelongsToOrgMixin, Base):
+class Resource(BelongsToOrgMixin, BelongsToRoomMixin, Base):
     """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
@@ -129,9 +137,6 @@ class Resource(BelongsToOrgMixin, Base):
     name: sa_orm.Mapped[str] = sa_orm.mapped_column(sa.Text)
     mime_type: sa_orm.Mapped[str] = sa_orm.mapped_column(sa.Text)
     url: sa_orm.Mapped[str] = sa_orm.mapped_column(sa.Text)
-    room_id: sa_orm.Mapped[RoomID] = sa_orm.mapped_column(
-        ForeignKey(Room).make(ondelete="CASCADE"), name="room_id"
-    )
     md5: sa_orm.Mapped[str] = sa_orm.mapped_column(sa.CHAR(32), nullable=True)
     size: sa_orm.Mapped[int] = sa_orm.mapped_column(nullable=True)
     original_path: sa_orm.Mapped[str] = sa_orm.mapped_column(nullable=True)
@@ -156,16 +161,13 @@ class Resource(BelongsToOrgMixin, Base):
     )
 
 
-class Source(BelongsToOrgMixin, Base):
+class Source(BelongsToOrgMixin, BelongsToRoomMixin, Base):
     """A source."""
 
     __tablename__ = "source"
     __table_args__ = (sa.UniqueConstraint("name", "room_id"),)
 
     name: sa_orm.Mapped[str] = sa_orm.mapped_column(sa.Text)
-    room_id: sa_orm.Mapped[RoomID] = sa_orm.mapped_column(
-        ForeignKey(Room).make(ondelete="CASCADE"),
-    )
     # 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()
@@ -196,7 +198,7 @@ class Source(BelongsToOrgMixin, Base):
         return self.name
 
 
-class Pipeline(BelongsToOrgMixin, Base):
+class Pipeline(BelongsToOrgMixin, BelongsToRoomMixin, Base):
     """A resource to source pipeline."""
 
     __tablename__ = "pipeline"
@@ -205,9 +207,6 @@ class Pipeline(BelongsToOrgMixin, Base):
     transformation: sa_orm.Mapped[pipeline_pb2.PipelineTransformation] = (
         sa_orm.mapped_column()
     )
-    room_id: sa_orm.Mapped[RoomID] = sa_orm.mapped_column(
-        ForeignKey(Room).make(ondelete="CASCADE")
-    )
     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()
@@ -227,7 +226,7 @@ class Pipeline(BelongsToOrgMixin, Base):
     )
 
 
-class PipelineInput(BelongsToOrgMixin, Base):
+class PipelineInput(BelongsToOrgMixin, BelongsToRoomMixin, Base):
     """Pipeline input resources."""
 
     __tablename__ = "pipeline_input"
@@ -250,7 +249,7 @@ class PipelineInput(BelongsToOrgMixin, Base):
     )
 
 
-class PipelineOutput(BelongsToOrgMixin, Base):
+class PipelineOutput(BelongsToOrgMixin, BelongsToRoomMixin, Base):
     """Objects for tracking pipeline output sources."""
 
     __tablename__ = "pipeline_output"
@@ -273,7 +272,7 @@ class PipelineOutput(BelongsToOrgMixin, Base):
     )
 
 
-class SourceResourceAssociation(BelongsToOrgMixin, Base):
+class SourceResourceAssociation(BelongsToOrgMixin, BelongsToRoomMixin, Base):
     __tablename__ = "source_resource_association"
 
     source_id: sa_orm.Mapped[SourceID | None] = (
@@ -294,7 +293,7 @@ class SourceResourceAssociation(BelongsToOrgMixin, Base):
     )
 
 
-class FeatureView(SoftDeleteMixin, BelongsToOrgMixin, Base):
+class FeatureView(SoftDeleteMixin, BelongsToOrgMixin, BelongsToRoomMixin, Base):
     """A FeatureView is a logical collection of sources used by various spaces."""
 
     __tablename__ = "feature_view"
@@ -304,12 +303,6 @@ class FeatureView(SoftDeleteMixin, BelongsToOrgMixin, Base):
     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="")
 
-    room_id: sa_orm.Mapped[RoomID | None] = sa_orm.mapped_column(
-        ForeignKey(Room).make(ondelete="CASCADE"),
-        nullable=True,
-        init=True,
-        default=None,
-    )
     room: sa_orm.Mapped[Room] = sa_orm.relationship(
         back_populates="feature_views", init=False
     )
@@ -333,10 +326,11 @@ class FeatureView(SoftDeleteMixin, BelongsToOrgMixin, Base):
     )
 
 
-class FeatureViewSource(BelongsToOrgMixin, Base):
+class FeatureViewSource(BelongsToOrgMixin, BelongsToRoomMixin, Base):
     """A source inside of a feature view."""
 
     __tablename__ = "feature_view_source"
+
     table_op_graph: sa_orm.Mapped[table_pb2.TableComputeOp] = sa_orm.mapped_column()
     id: sa_orm.Mapped[FeatureViewSourceID | None] = primary_key_identity_column()
     drop_disconnected: sa_orm.Mapped[bool] = sa_orm.mapped_column(default=False)
@@ -356,18 +350,12 @@ class FeatureViewSource(BelongsToOrgMixin, Base):
     )
 
 
-class Space(BelongsToOrgMixin, Base):
+class Space(BelongsToOrgMixin, BelongsToRoomMixin, Base):
     """A space is a named evaluation of space parameters."""
 
     __tablename__ = "space"
     __table_args__ = (sa.UniqueConstraint("name", "room_id"),)
 
-    room_id: sa_orm.Mapped[RoomID] = sa_orm.mapped_column(
-        ForeignKey(Room).make(ondelete="CASCADE"),
-        nullable=True,
-        init=True,
-        default=None,
-    )
     room: sa_orm.Mapped[Room] = sa_orm.relationship(
         back_populates="spaces", init=True, default=None
     )
@@ -404,7 +392,7 @@ class Space(BelongsToOrgMixin, Base):
         return self.name
 
 
-class SpaceRun(BelongsToOrgMixin, Base):
+class SpaceRun(BelongsToOrgMixin, BelongsToRoomMixin, Base):
     """A Space run."""
 
     __tablename__ = "space_run"
@@ -439,18 +427,12 @@ class SpaceRun(BelongsToOrgMixin, Base):
     )
 
 
-class Agent(SoftDeleteMixin, BelongsToOrgMixin, Base):
+class Agent(SoftDeleteMixin, BelongsToOrgMixin, BelongsToRoomMixin, Base):
     """An Agent."""
 
     __tablename__ = "agent"
     __table_args__ = (live_unique_constraint("name", "room_id"),)
 
-    room_id: sa_orm.Mapped[RoomID | None] = sa_orm.mapped_column(
-        ForeignKey(Room).make(ondelete="CASCADE"),
-        nullable=True,
-        init=True,
-        default=None,
-    )
     id: sa_orm.Mapped[AgentID | None] = primary_key_identity_column()
     name: sa_orm.Mapped[str] = sa_orm.mapped_column(sa.Text, default=None)
 
@@ -478,7 +460,7 @@ class Agent(SoftDeleteMixin, BelongsToOrgMixin, Base):
     )
 
 
-class AgentSpaceAssociation(BelongsToOrgMixin, Base):
+class AgentSpaceAssociation(BelongsToOrgMixin, BelongsToRoomMixin, Base):
     __tablename__ = "agent_space_association"
 
     agent_id: sa_orm.Mapped[AgentID] = primary_key_foreign_column(
@@ -495,7 +477,7 @@ class AgentSpaceAssociation(BelongsToOrgMixin, Base):
     )
 
 
-class UserMessage(SoftDeleteMixin, BelongsToOrgMixin, Base):
+class UserMessage(SoftDeleteMixin, BelongsToOrgMixin, BelongsToRoomMixin, Base):
     """A message sent by an user."""
 
     __tablename__ = "user_message"
@@ -509,7 +491,7 @@ class UserMessage(SoftDeleteMixin, BelongsToOrgMixin, Base):
     message: sa_orm.Mapped[str | None] = sa_orm.mapped_column(sa.Text, default=None)
 
 
-class AgentMessage(SoftDeleteMixin, BelongsToOrgMixin, Base):
+class AgentMessage(SoftDeleteMixin, BelongsToOrgMixin, BelongsToRoomMixin, Base):
     """A message sent by an agent."""
 
     __tablename__ = "agent_message"
@@ -537,7 +519,7 @@ class AgentMessage(SoftDeleteMixin, BelongsToOrgMixin, Base):
     )
 
 
-class MessageEntry(SoftDeleteMixin, BelongsToOrgMixin, Base):
+class MessageEntry(SoftDeleteMixin, BelongsToOrgMixin, BelongsToRoomMixin, Base):
     """A message either sent by an Agent or an User."""
 
     __tablename__ = "message_entry"
@@ -580,8 +562,9 @@ class CompletionModel(SoftDeleteMixin, BelongsToOrgMixin, Base):
     id: sa_orm.Mapped[CompletionModelID | None] = 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)
-    model_name: sa_orm.Mapped[str] = sa_orm.mapped_column(sa.Text, default=None)
-    endpoint: sa_orm.Mapped[str] = sa_orm.mapped_column(sa.Text, default=None)
+    parameters: sa_orm.Mapped[completion_model_pb2.CompletionModelParameters | None] = (
+        sa_orm.mapped_column(default=None)
+    )
     secret_api_key: sa_orm.Mapped[str] = sa_orm.mapped_column(sa.Text, default=None)
 
     @property

corvic/orm/base.py

@@ -41,6 +41,7 @@ from corvic.orm.keys import (
 from corvic_generated.orm.v1 import (
     agent_pb2,
     common_pb2,
+    completion_model_pb2,
     feature_view_pb2,
     pipeline_pb2,
     space_pb2,
@@ -112,6 +113,9 @@ class Base(sa_orm.MappedAsDataclass, sa_orm.DeclarativeBase):
             pipeline_pb2.PipelineTransformation()
         ),
         event_pb2.Event: ProtoMessageDecorator(event_pb2.Event()),
+        completion_model_pb2.CompletionModelParameters: ProtoMessageDecorator(
+            completion_model_pb2.CompletionModelParameters()
+        ),
         # ID types
         OrgID: StrIDDecorator(OrgID()),
         RoomID: IntIDDecorator(RoomID()),