corvic-engine 0.3.0rc46__tar.gz → 0.3.0rc47__tar.gz

{corvic_engine-0.3.0rc46 → corvic_engine-0.3.0rc47}/Cargo.lock

@@ -22,7 +22,7 @@ checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd"
 
 [[package]]
 name = "corvic-engine"
-version = "0.3.0-rc.46"
+version = "0.3.0-rc.47"
 dependencies = [
  "numpy",
  "pyo3",
@@ -63,7 +63,7 @@ checksum = "60b1af1c220855b6ceac025d3f6ecdd2b7c4894bfe9cd9bda4fbb4bc7c0d4cf0"
 
 [[package]]
 name = "engine-walk"
-version = "0.3.0-rc.46"
+version = "0.3.0-rc.47"
 dependencies = [
  "version_check",
 ]
@@ -335,9 +335,9 @@ checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49"
 
 [[package]]
 name = "smallvec"
-version = "1.13.2"
+version = "1.14.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "3c5e1a9a646d36c3599cd173a41282daf47c44583ad367b8e6837255952e5c67"
+checksum = "7fcf8323ef1faaee30a44a340193b1ac6814fd9b7b4e88e9d4519a3e4abe1cfd"
 
 [[package]]
 name = "syn"
@@ -368,9 +368,9 @@ dependencies = [
 
 [[package]]
 name = "unicode-ident"
-version = "1.0.16"
+version = "1.0.17"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "a210d160f08b701c8721ba1c726c11662f877ea6b7094007e1ca9a1041945034"
+checksum = "00e2473a93778eb0bad35909dff6a10d28e63f792f16ed15e404fca9d5eeedbe"
 
 [[package]]
 name = "unindent"

{corvic_engine-0.3.0rc46 → corvic_engine-0.3.0rc47}/Cargo.toml

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

{corvic_engine-0.3.0rc46 → corvic_engine-0.3.0rc47}/PKG-INFO

@@ -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 → corvic_engine-0.3.0rc47}/pyproject.toml

@@ -20,7 +20,7 @@ authors = [
 ]
 readme = "README.md"
 description = "Seamless embedding generation and retrieval."
-version = "0.3.0-rc.46" # sync this with tool.poetry.version below
+version = "0.3.0-rc.47" # 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.
@@ -74,7 +74,7 @@ packages = [
 ]
 description = ""
 authors = []
-version = "0.3.0-rc.46"
+version = "0.3.0-rc.47"
 
 [tool.poetry.scripts]
 check = "corvic_check.cli:main"

{corvic_engine-0.3.0rc46 → corvic_engine-0.3.0rc47}/python/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_engine-0.3.0rc46 → corvic_engine-0.3.0rc47}/python/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_engine-0.3.0rc46 → corvic_engine-0.3.0rc47}/python/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_engine-0.3.0rc46 → corvic_engine-0.3.0rc47}/python/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_engine-0.3.0rc46 → corvic_engine-0.3.0rc47}/python/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_engine-0.3.0rc46 → corvic_engine-0.3.0rc47}/python/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: ...