corvic-engine 0.3.0rc46__cp38-abi3-win_amd64.whl → 0.3.0rc47__cp38-abi3-win_amd64.whl

corvic/model/_agent.py

@@ -3,7 +3,8 @@
 from __future__ import annotations
 
 import copy
-from collections.abc import Sequence
+import datetime
+from collections.abc import Iterable, Sequence
 from typing import TypeAlias
 
 from sqlalchemy import orm as sa_orm
@@ -95,6 +96,39 @@ class Agent(BaseModel[AgentID, models_pb2.Agent, orm.Agent]):
             ),
         )
 
+    @classmethod
+    def from_proto(
+        cls, proto: models_pb2.Agent, client: system.Client | None = None
+    ) -> Agent:
+        client = client or Defaults.get_default_client()
+        return cls(client, proto)
+
+    @classmethod
+    def list(
+        cls,
+        *,
+        limit: int | None = None,
+        room_id: RoomID | None = None,
+        created_before: datetime.datetime | None = None,
+        client: system.Client | None = None,
+        ids: Iterable[AgentID] | None = None,
+        existing_session: sa_orm.Session | None = None,
+    ) -> Ok[list[Agent]] | NotFoundError | InvalidArgumentError:
+        """List agent models."""
+        client = client or Defaults.get_default_client()
+        match cls.list_as_proto(
+            client,
+            limit=limit,
+            room_id=room_id,
+            created_before=created_before,
+            ids=ids,
+            existing_session=existing_session,
+        ):
+            case NotFoundError() | InvalidArgumentError() as err:
+                return err
+            case Ok(protos):
+                return Ok([cls.from_proto(proto, client) for proto in protos])
+
     @property
     def name(self) -> str:
         return self.proto_self.name

corvic/model/_base_model.py

@@ -206,7 +206,8 @@ class BaseModel(Generic[ID, _ProtoObj, OrmObj], UsesOrmID[ID, _ProtoObj]):
         try:
             _ = self.proto_to_orm(self.proto_self, session).unwrap_or_raise()
             session.flush()
-        except sa.exc.DatabaseError as err:
+        # TODO(thunt): Possibly separate out DatabaseError into a precondition error
+        except (sa.exc.DatabaseError, sa.exc.StatementError) as err:
             return InvalidArgumentError.from_(err)
         return Ok(None)
 

corvic/model/_proto_orm_convert.py

@@ -386,6 +386,37 @@ def source_proto_to_orm(
     return Ok(_add_orm_to_session(orm_obj, proto_obj.org_id, session))
 
 
+def _update_agent_associations(
+    agent: orm.Agent,
+    session: sa_orm.Session,
+) -> list[orm.SpaceID]:
+    associated_space_ids = list[orm.SpaceID]()
+    agent_id = agent.id
+    agent_parameters = agent.parameters
+    if not agent_id or not agent_parameters:
+        return associated_space_ids
+    for (
+        raw_space_id,
+        instruction,
+    ) in agent_parameters.orchestrator_parameters.space_instructions.items():
+        space_id = orm.SpaceID(raw_space_id)
+        space_run_id = (
+            orm.SpaceRunID(instruction.space_run_id)
+            if instruction.space_run_id
+            else None
+        )
+        associated_space_ids.append(space_id)
+        session.merge(
+            orm.AgentSpaceAssociation(
+                room_id=agent.room_id,
+                agent_id=agent_id,
+                space_id=space_id,
+                space_run_id=space_run_id,
+            )
+        )
+    return associated_space_ids
+
+
 def agent_proto_to_orm(
     proto_obj: models_pb2.Agent, session: sa_orm.Session
 ) -> Ok[orm.Agent] | orm.InvalidORMIdentifierError | InvalidArgumentError:
@@ -403,7 +434,18 @@ def agent_proto_to_orm(
         parameters=proto_obj.agent_parameters,
         room_id=room_id,
     )
-    return Ok(_add_orm_to_session(orm_obj, proto_obj.org_id, session))
+    agent = _add_orm_to_session(orm_obj, proto_obj.org_id, session)
+    session.flush()
+
+    if not agent.id:
+        return InvalidArgumentError("failed to add agent to session")
+    associated_space_ids = _update_agent_associations(agent, session)
+    session.execute(
+        sa.delete(orm.AgentSpaceAssociation)
+        .where(orm.AgentSpaceAssociation.agent_id == agent.id)
+        .where(orm.AgentSpaceAssociation.space_id.not_in(associated_space_ids))
+    )
+    return Ok(agent)
 
 
 def space_proto_to_orm(
@@ -650,6 +692,11 @@ def resource_delete_orms(
 def agent_delete_orms(
     ids: Sequence[orm.AgentID], session: orm.Session
 ) -> Ok[None] | InvalidArgumentError:
+    session.execute(
+        sa.delete(orm.AgentSpaceAssociation).where(
+            orm.AgentSpaceAssociation.agent_id.in_(ids)
+        )
+    )
     session.execute(sa.delete(orm.Agent).where(orm.Agent.id.in_(ids)))
     return Ok(None)
 
@@ -718,12 +765,14 @@ def room_delete_orms(
     ids: Sequence[orm.RoomID], session: orm.Session
 ) -> Ok[None] | InvalidArgumentError:
     source_refs = list(
-        session.scalars(sa.select(orm.Source).where(orm.Source.room_id == ids).limit(1))
+        session.scalars(
+            sa.select(orm.Source).where(orm.Source.room_id.in_(ids)).limit(1)
+        )
     )
     if source_refs:
         return InvalidArgumentError("cannot delete a room that still has sources")
 
-    session.execute(sa.delete(orm.Room).where(orm.Room.id == ids))
+    session.execute(sa.delete(orm.Room).where(orm.Room.id.in_(ids)))
     return Ok(None)
 
 

corvic/orm/__init__.py

@@ -286,14 +286,6 @@ class Space(BelongsToOrgMixin, BelongsToRoomMixin, Base):
         viewonly=True,
     )
 
-    agent_associations: sa_orm.Mapped[list[AgentSpaceAssociation]] = (
-        sa_orm.relationship(
-            init=False,
-            viewonly=True,
-            default_factory=list,
-        )
-    )
-
     @property
     def space_key(self):
         return self.name
@@ -359,15 +351,15 @@ class Agent(SoftDeleteMixin, BelongsToOrgMixin, BelongsToRoomMixin, Base):
 class AgentSpaceAssociation(BelongsToOrgMixin, BelongsToRoomMixin, Base):
     __tablename__ = "agent_space_association"
 
+    space_run_id: sa_orm.Mapped[SpaceRunID | None] = sa_orm.mapped_column(
+        ForeignKey(SpaceRun).make()
+    )
     agent_id: sa_orm.Mapped[AgentID] = primary_key_foreign_column(
-        ForeignKey(Agent).make()
+        ForeignKey(Agent).make(ondelete="CASCADE")
     )
     space_id: sa_orm.Mapped[SpaceID] = primary_key_foreign_column(
         ForeignKey(Space).make()
     )
-    space: sa_orm.Mapped[Space] = sa_orm.relationship(
-        back_populates="agent_associations", init=False, viewonly=True
-    )
 
 
 class CompletionModel(SoftDeleteMixin, BelongsToOrgMixin, Base):

corvic/pa_scalar/_to_value.py

@@ -1,11 +1,13 @@
 import base64
 import datetime
 import decimal
+import math
 from collections.abc import Callable, Iterable, Mapping
-from typing import Final, TypeAlias, cast
+from typing import Any, Final, TypeAlias, cast
 
 import pyarrow as pa
 from google.protobuf import struct_pb2
+from pyarrow.lib import ArrowException
 
 from corvic.pa_scalar._temporal import (
     MonthDayNanoFrac,
@@ -118,11 +120,17 @@ def _visit_uint64(
     return struct_pb2.Value(string_value=str(value))
 
 
+def _is_invalid_float(value: float):
+    return math.isnan(value) or math.isinf(value)
+
+
 def _visit_float16(
     value: _PyValue,
     dtype: pa.DataType,
 ) -> struct_pb2.Value:
     value = cast(float, value)
+    if _is_invalid_float(value):
+        return _visit_null(value, pa.null())
     return struct_pb2.Value(number_value=value)
 
 
@@ -131,6 +139,8 @@ def _visit_float32(
     dtype: pa.DataType,
 ) -> struct_pb2.Value:
     value = cast(float, value)
+    if _is_invalid_float(value):
+        return _visit_null(value, pa.null())
     return struct_pb2.Value(number_value=value)
 
 
@@ -139,6 +149,8 @@ def _visit_float64(
     dtype: pa.DataType,
 ) -> struct_pb2.Value:
     value = cast(float, value)
+    if _is_invalid_float(value):
+        return _visit_null(value, pa.null())
     return struct_pb2.Value(number_value=value)
 
 
@@ -412,6 +424,8 @@ def to_value(
         case int():
             return struct_pb2.Value(string_value=str(scalar))
         case float():
+            if _is_invalid_float(scalar):
+                return to_value(None)
             return struct_pb2.Value(number_value=scalar)
         case str():
             return struct_pb2.Value(string_value=scalar)
@@ -424,6 +438,36 @@ def to_value(
             return _visit(value, scalar.type)
 
 
+def _to_py(scalar: Scalar) -> Any:
+    try:
+        # This scalar might have malformed binary data that's incorrectly typed
+        # as string (e.g. non-utf-8 strings), so skip it
+        scalar.validate(full=True)
+        return scalar.as_py()
+    except ArrowException:
+        # To create examples that hit this code path, one needs to create a
+        # pyarrow.BinaryArray and then cast it to a StringArray but also supply
+        # a CastOptions with allow_invalid_utf8 set to True.
+        # TODO(aneesh): also check pa.types.is_dictionary
+        if (
+            pa.types.is_list(scalar.type)
+            or pa.types.is_large_list(scalar.type)
+            or pa.types.is_fixed_size_list(scalar.type)
+        ):
+            lscalar = cast(pa.ListScalar[Any], scalar)
+            return [_to_py(v) for v in lscalar.values]  # noqa: PD011 # type: ignore[unknownMemberType]
+        if pa.types.is_map(scalar.type):
+            mscalar = cast(pa.MapScalar[Any, Any], scalar)
+            return {
+                res["key"]: res["value"]
+                for res in [_to_py(v) for v in mscalar.values]  # noqa: PD011 # type: ignore[unknownMemberType]
+            }
+        if pa.types.is_struct(scalar.type):
+            sscalar = cast(pa.StructScalar, scalar)
+            return {k: _to_py(sscalar[k]) for k in sscalar}
+        return None
+
+
 def batch_to_structs(
     batch: pa.RecordBatch,
 ) -> list[struct_pb2.Struct]:
@@ -432,7 +476,8 @@ def batch_to_structs(
     field_by_name = {field.name: field for field in schema}
 
     ret: list[struct_pb2.Struct] = []
-    for row in batch.to_pylist():
+    for pa_row in batch.to_struct_array():
+        row = _to_py(cast(Scalar, pa_row))
         fields = {
             col_name: _visit(v, field_by_name[col_name].type)
             for col_name, v in row.items()

corvic/pa_scalar/_types.py

@@ -39,3 +39,5 @@ class Scalar(Protocol):
     def equals(self, other: Any) -> bool: ...
 
     def as_py(self) -> Any: ...
+
+    def validate(self, *, full: bool = False) -> bool | None: ...

{corvic_engine-0.3.0rc46.dist-info → corvic_engine-0.3.0rc47.dist-info}/METADATA

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

{corvic_engine-0.3.0rc46.dist-info → corvic_engine-0.3.0rc47.dist-info}/RECORD

@@ -1,6 +1,6 @@
-corvic_engine-0.3.0rc46.dist-info/METADATA,sha256=_LLt2PwaSBy5Jy3CuT9BHywURG_SkNOW0g4yJ9yahC8,1876
-corvic_engine-0.3.0rc46.dist-info/WHEEL,sha256=_g1M2QM3kt1Ssm_sHOg_3TUY7GxNE2Ueyslb9ZDtPwk,94
-corvic_engine-0.3.0rc46.dist-info/licenses/LICENSE,sha256=DSS1OD0oIgssKOmAzkMRBv5jvvVuZQbrIv8lpl9DXY8,1035
+corvic_engine-0.3.0rc47.dist-info/METADATA,sha256=DACl2w5jIjPNPHgd9k7VlUxQUC0fi2pN3UewAPoY4f8,1876
+corvic_engine-0.3.0rc47.dist-info/WHEEL,sha256=_g1M2QM3kt1Ssm_sHOg_3TUY7GxNE2Ueyslb9ZDtPwk,94
+corvic_engine-0.3.0rc47.dist-info/licenses/LICENSE,sha256=DSS1OD0oIgssKOmAzkMRBv5jvvVuZQbrIv8lpl9DXY8,1035
 corvic/context/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 corvic/context/__init__.py,sha256=zBnPiP-tStGSVMG_0-G_0ay6-yIX2aerW_oYRzAex74,1702
 corvic/embed/node2vec.py,sha256=JnYb8f2g4XhF6LL2TjpMxLfKhn_Yp1AzptsWwrKQWgc,11146
@@ -13,15 +13,15 @@ corvic/engine/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 corvic/engine/_native.pyi,sha256=KYMPtvXqHZ-jMgZohLf4se3rr-rBpCihmjANcr6s8ag,1390
 corvic/engine/__init__.py,sha256=XL4Vg7rNcBi29ccVelpeFizR9oJtGYXDn84W9zok9d4,975
 corvic/model/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-corvic/model/_agent.py,sha256=WGdxu0oLHU7EEgBOMQ0UOnH2AIIUGfeH6VSQZgC0kEk,3621
-corvic/model/_base_model.py,sha256=m4JkZu8QAoQEeHqzpMUlH08PGOjVZbqld4Rm0uW73QQ,8225
+corvic/model/_agent.py,sha256=8qmi1x1pBQCXgz7Zr2YdrWRHRZkbp63cyJw7j9OvqQY,4818
+corvic/model/_base_model.py,sha256=qRTZBBHdkqMH6xqX1FwtTm4cWFC4aPDktJeGyxyxtpk,8336
 corvic/model/_completion_model.py,sha256=uoqF7hwxzGXXqSPZT_CIcNBSDmYhcxMotpGPucWH6Q0,6656
 corvic/model/_defaults.py,sha256=yoKPPSmYJCE5YAD5jLTEmT4XNf_zXoggNK-uyG8MfVs,1524
 corvic/model/_errors.py,sha256=Ctlq04SDwHzJPvLaL1rzqzwVqf2b50EILfW3cH4vnh8,261
 corvic/model/_feature_type.py,sha256=Y-_-wa9fv7XaCAkxfjjoCLxxK2Ftfba-PMefD7bNXzs,917
 corvic/model/_feature_view.py,sha256=kM8kMzb1oxAgIOMVmv6twZ-JLeP2NopdsKHCiKjuNIE,49700
 corvic/model/_pipeline.py,sha256=A_q_nWm6UBN-AKlbQkhWNMG2r-uW0IR6vGJbhYv7z3k,17578
-corvic/model/_proto_orm_convert.py,sha256=6CWe_b-L7JBS4Pb_SjUeuiuczWh-nPw0jElQIMOl_UA,25029
+corvic/model/_proto_orm_convert.py,sha256=rq4SJYsTCHyOkZswI9FYfj1Ne9Heed117ii89ChjPvc,26607
 corvic/model/_resource.py,sha256=O93gy2r4HjsZzVUahrGbLGmxI7sfSVQCvwvCLN3nSXQ,7097
 corvic/model/_room.py,sha256=57MiBfj8hZcmUfq2PeECrOWDpBZAOSjnVqNUIGXOy2Q,2898
 corvic/model/_source.py,sha256=JBCk1I6u_rUKPiB4Fvtl7uVm0Jx0LF1oWNd1-Wn_sbI,9412
@@ -50,13 +50,13 @@ corvic/orm/keys.py,sha256=Ag6Xbpvxev-VByT1KJ8ChUn9vKVEzkkMXxrjvtADCtY,2182
 corvic/orm/mixins.py,sha256=HfmzJ7LblHtddbbkDmv7nNWURL87Bnj8NeOnNbfmSN4,17794
 corvic/orm/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 corvic/orm/_proto_columns.py,sha256=tcOu92UjFJFYZLasS6sWJQBDRK26yrnmpTii_LDY4iw,913
-corvic/orm/__init__.py,sha256=3HLMICNB1iMyMS70tlvaap8WDrAfxqSxp7W1O-cnHc4,14656
+corvic/orm/__init__.py,sha256=rOYy3hi3fVbAkXIoJ5NZLDz3uHdw3Bko2ZpRDtb3Fkg,14449
 corvic/pa_scalar/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 corvic/pa_scalar/_const.py,sha256=1nk6w3Y7crd3J5jSCq7DRVa1lcGk4H1RUr1l4NjnlzE,868
 corvic/pa_scalar/_from_value.py,sha256=fS3TNPcPI3jAKGmcUIhn8rdqdQEAwgTLEneVxFUeK6M,27531
 corvic/pa_scalar/_temporal.py,sha256=HPL1z5k4frIiGWM469YGj9j0hCPXNHtI_X1jk1y1c1Q,7842
-corvic/pa_scalar/_to_value.py,sha256=F11uEikSft9AahBqA7aHmvHObwsTqoZN90KC4ucbem0,11574
-corvic/pa_scalar/_types.py,sha256=73dusgldts4o8j-FgoKjER9malSlMbGVrD4hP3mflWE,1062
+corvic/pa_scalar/_to_value.py,sha256=C1l-BVsS4Rs_KL8hVPllpzjP0NVVwfaOb0WM2pjZhoA,13476
+corvic/pa_scalar/_types.py,sha256=shbytO0ji-H2rBOX_1fooVOshb22wwkVU1W99VBKz1A,1131
 corvic/pa_scalar/__init__.py,sha256=1nfc0MFGpw78RQEI13VE5hpHuyw_DoE7sJbmzqx5pws,1063
 corvic/proto_wrapper/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 corvic/proto_wrapper/_errors.py,sha256=0HFmBK9EGYi4lXJ3RFsktHoPZn2m3cMKroI0Eir0JcM,114
@@ -244,5 +244,5 @@ corvic_generated/status/v1/event_pb2.pyi,sha256=eU-ibrYpvEAJSIDlSa62-bC96AQU1ykF
 corvic_generated/status/v1/event_pb2_grpc.pyi,sha256=H9-ADaiKR9iyVZvmnXutZqWwRRCDxjUIktkfJrJFIHg,417
 corvic_generated/status/v1/service_pb2.pyi,sha256=iXLR2FOKQJpBgvBzpD2kVwcYOCksP2aRwK4JYaI9CBw,558
 corvic_generated/status/v1/service_pb2_grpc.pyi,sha256=OoAnaZ64FD0UTzPoRhYvQU8ecoilhHj3ySjSfHbVDaU,1501
-corvic/engine/_native.pyd,sha256=6ts0z2Hw7EtkPpwjGCBvZqq9Pg282NoTZk5EwiiLZtc,438272
-corvic_engine-0.3.0rc46.dist-info/RECORD,,
+corvic/engine/_native.pyd,sha256=mlFXn3uT-o8-APpr9ihxwWATcd-mdiSGtZvkGcoJibs,438272
+corvic_engine-0.3.0rc47.dist-info/RECORD,,