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

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

@@ -22,7 +22,7 @@ checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd"
 
 [[package]]
 name = "corvic-engine"
-version = "0.3.0-rc.47"
+version = "0.3.0-rc.49"
 dependencies = [
  "numpy",
  "pyo3",
@@ -63,7 +63,7 @@ checksum = "60b1af1c220855b6ceac025d3f6ecdd2b7c4894bfe9cd9bda4fbb4bc7c0d4cf0"
 
 [[package]]
 name = "engine-walk"
-version = "0.3.0-rc.47"
+version = "0.3.0-rc.49"
 dependencies = [
  "version_check",
 ]

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

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

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

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

{corvic_engine-0.3.0rc47 → corvic_engine-0.3.0rc49}/pyproject.toml

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

{corvic_engine-0.3.0rc47 → corvic_engine-0.3.0rc49}/python/corvic/model/_feature_view.py

@@ -340,7 +340,9 @@ class Relationship:
 
         result = self.joined_table().select(result_columns)
 
-        for batch in result.to_polars().unwrap_or_raise():
+        for batch in result.to_polars(
+            room_id=self.start_source.room_id
+        ).unwrap_or_raise():
             for row in batch.rows(named=True):
                 yield (row[result_columns[0]], row[result_columns[1]])
 

{corvic_engine-0.3.0rc47 → corvic_engine-0.3.0rc49}/python/corvic/model/_resource.py

@@ -4,7 +4,6 @@ from __future__ import annotations
 
 import copy
 import datetime
-import uuid
 from collections.abc import Iterable, Sequence
 from typing import TypeAlias
 
@@ -219,8 +218,8 @@ class Resource(BaseModel[ResourceID, models_pb2.Resource, orm.Resource]):
     ) -> Self:
         client = client or Defaults.get_default_client()
         room_id = room_id or Defaults.get_default_room_id(client)
-        blob = client.storage_manager.tabular.blob(f"polars_dataframe/{uuid.uuid4()}")
 
+        blob = client.storage_manager.make_tabular_blob(room_id, "anonymous_tables")
         with blob.open(mode="wb") as stream:
             data_frame.write_parquet(stream)
 

{corvic_engine-0.3.0rc47 → corvic_engine-0.3.0rc49}/python/corvic/system/__init__.py

@@ -48,7 +48,6 @@ from corvic.system.storage import (
     Blob,
     BlobClient,
     Bucket,
-    DataKindManager,
     DataMisplacedError,
     StorageManager,
 )
@@ -71,7 +70,6 @@ __all__ = [
     "Clip",
     "ClipText",
     "DEFAULT_VECTOR_COLUMN_NAMES_TO_SIZES",
-    "DataKindManager",
     "DataMisplacedError",
     "DimensionReducer",
     "EmbedTextContext",

{corvic_engine-0.3.0rc47 → corvic_engine-0.3.0rc49}/python/corvic/system/storage.py

@@ -3,14 +3,14 @@
 from __future__ import annotations
 
 import contextlib
-import io
 import uuid
 from collections.abc import Iterator
-from typing import Any, Literal
+from typing import Any, BinaryIO, Final, Literal
 
 from typing_extensions import Protocol
 
-from corvic.result import Error, Ok
+from corvic import orm
+from corvic.result import Error
 
 
 class DataMisplacedError(Error):
@@ -74,9 +74,7 @@ class Blob(Protocol):
     def url(self) -> str: ...
 
     @contextlib.contextmanager
-    def open(
-        self, mode: Literal["rb", "wb"], **kwargs: Any
-    ) -> Iterator[io.BytesIO]: ...
+    def open(self, mode: Literal["rb", "wb"], **kwargs: Any) -> Iterator[BinaryIO]: ...
 
     def delete(self) -> None: ...
 
@@ -105,60 +103,16 @@ class Bucket(Protocol):
     def list_blobs(self, prefix: str | None = None) -> Iterator[Blob]: ...
 
 
-class DataKindManager:
-    """Manages the names of blobs that corvic stores for a particular data kind.
-
-    Kinds are managed by the NamespaceManager
-    """
-
-    def __init__(self, storage_manager: StorageManager, prefix: str):
-        if prefix.endswith("/"):
-            raise ValueError("prefix should not end with a path separator (/)")
-        self._namespace_manager = storage_manager
-        self._prefix = prefix
-
-    @property
-    def prefix(self):
-        return self._prefix
-
-    def make_anonymous_table_url(self):
-        return self.blob(f"anonymous_tables/{uuid.uuid4()}.parquet").url
-
-    def blob(self, blob_name: str):
-        return self._namespace_manager.bucket.blob(f"{self.prefix}/{blob_name}")
-
-    def _blob_from_url(self, url: str) -> Ok[Blob] | DataMisplacedError:
-        blob = self._namespace_manager.blob_from_url(url)
-        if self._namespace_manager.bucket.name != blob.bucket.name:
-            return DataMisplacedError(
-                "data stored at a different bucket than expected",
-                url=url,
-                expected_bucket=self._namespace_manager.bucket.name,
-            )
-        if not blob.name.startswith(self._prefix):
-            return DataMisplacedError(
-                "data stored at a different prefix than expected",
-                url=url,
-                expected_prefix=self._prefix,
-            )
-        return Ok(blob)
-
-    def blob_name_from_url(self, url: str) -> Ok[str] | DataMisplacedError:
-        def _to_name(blob: Blob) -> str:
-            return blob.name.removeprefix(self._prefix + "/")
-
-        return self._blob_from_url(url=url).map(_to_name)
-
-
 class StorageManager:
     """Manages the names of blobs that corvic stores."""
 
     _blob_client: BlobClient
     _bucket_name: str
-    _unstructured_manager: DataKindManager
-    _tabular_manager: DataKindManager
-    _space_run_manager: DataKindManager
-    _vector_manager: DataKindManager
+
+    unstructured_prefix: Final[str]
+    tabular_prefix: Final[str]
+    vector_prefix: Final[str]
+    bucket: Final[Bucket]
 
     def __init__(
         self,
@@ -167,36 +121,55 @@ class StorageManager:
         bucket_name: str,
         unstructured_prefix: str,
         tabular_prefix: str,
-        space_run_prefix: str,
         vector_prefix: str,
     ):
         self._blob_client = blob_client
         self._bucket_name = bucket_name
 
-        self._unstructured_manager = DataKindManager(self, unstructured_prefix)
-        self._tabular_manager = DataKindManager(self, tabular_prefix)
-        self._space_run_manager = DataKindManager(self, space_run_prefix)
-        self._vector_manager = DataKindManager(self, vector_prefix)
-
-    def blob_from_url(self, url: str):
-        return self._blob_client.blob_from_url(url)
+        if unstructured_prefix.endswith("/") or tabular_prefix.endswith("/"):
+            raise ValueError("prefix should not end with a path separator (/)")
 
-    @property
-    def bucket(self) -> Bucket:
-        return self._blob_client.bucket(self._bucket_name)
+        self.unstructured_prefix = unstructured_prefix
+        self.tabular_prefix = tabular_prefix
+        self.vector_prefix = vector_prefix
 
-    @property
-    def tabular(self):
-        return self._tabular_manager
+        self.bucket = self._blob_client.bucket(self._bucket_name)
 
-    @property
-    def unstructured(self):
-        return self._unstructured_manager
+    @staticmethod
+    def _render_room_id(room_id: orm.RoomID) -> str:
+        return f"{int(str(room_id)):016}"
 
-    @property
-    def space_run(self):
-        return self._space_run_manager
+    def blob_from_url(self, url: str):
+        return self._blob_client.blob_from_url(url)
 
-    @property
-    def vector(self):
-        return self._vector_manager
+    def make_tabular_blob(self, room_id: orm.RoomID, suffix: str | None = None) -> Blob:
+        if suffix:
+            name = f"{self._render_room_id(room_id)}/{uuid.uuid4()}/{suffix}"
+        else:
+            name = f"{self._render_room_id(room_id)}/{uuid.uuid4()}"
+        return self.bucket.blob(f"{self.tabular_prefix}/{name}")
+
+    def make_unstructured_blob(
+        self, room_id: orm.RoomID, suffix: str | None = None
+    ) -> Blob:
+        if suffix:
+            name = f"{self._render_room_id(room_id)}/{uuid.uuid4()}/{suffix}"
+        else:
+            name = f"{self._render_room_id(room_id)}/{uuid.uuid4()}"
+        return self.bucket.blob(f"{self.unstructured_prefix}/{name}")
+
+    def make_vector_blob(self, room_id: orm.RoomID, suffix: str | None = None) -> Blob:
+        if suffix:
+            name = f"{self._render_room_id(room_id)}/{uuid.uuid4()}/{suffix}"
+        else:
+            name = f"{self._render_room_id(room_id)}/{uuid.uuid4()}"
+        return self.bucket.blob(f"{self.vector_prefix}/{name}")
+
+    def get_tabular_blob_from_blob_name(self, blob_name: str) -> Blob:
+        return self.bucket.blob(f"{self.tabular_prefix}/{blob_name}")
+
+    def get_unstructured_blob_from_blob_name(self, blob_name: str) -> Blob:
+        return self.bucket.blob(f"{self.unstructured_prefix}/{blob_name}")
+
+    def get_vector_blob_from_blob_name(self, blob_name: str) -> Blob:
+        return self.bucket.blob(f"{self.vector_prefix}/{blob_name}")

{corvic_engine-0.3.0rc47 → corvic_engine-0.3.0rc49}/python/corvic/system_sqlite/client.py

@@ -135,7 +135,6 @@ class Client(corvic.system.Client):
             bucket_name=bucket.name,
             unstructured_prefix="unstructured_data",
             tabular_prefix="tabular_data",
-            space_run_prefix="experiment_run_data",
             vector_prefix="vectors",
         )
         corvic.orm.Base.metadata.create_all(self._sa_engine)

{corvic_engine-0.3.0rc47 → corvic_engine-0.3.0rc49}/python/corvic/system_sqlite/staging.py

@@ -15,7 +15,6 @@ import sqlglot.expressions
 from corvic.op_graph import Schema
 from corvic.result import InternalError, Ok
 from corvic.system import (
-    DataKindManager,
     StagingDB,
     StorageManager,
     TableSliceArgs,
@@ -255,17 +254,18 @@ class DuckDBStaging(StagingDB):
 
     def _update_blobs_for_prefix(
         self,
-        manager: DataKindManager,
         next_count: dict[str, int | None],
+        prefix: str,
         transform_table: Callable[[str, pa.Table], None] | None = None,
     ):
         bucket = self._storage_manager.bucket
         blobs = bucket.list_blobs()
-        prefix = manager.prefix
         table_blobs = [
             (
                 blob,
-                manager.blob_name_from_url(url=blob.url).unwrap_or_raise(),
+                self._storage_manager.blob_from_url(blob.url).name.removeprefix(
+                    prefix + "/"
+                ),
             )
             for blob in blobs
             if blob.name.startswith(prefix)
@@ -289,9 +289,14 @@ class DuckDBStaging(StagingDB):
         As a side-effect, update _table_counts.
         """
         next_count: dict[str, int | None] = {}
-        self._update_blobs_for_prefix(self._storage_manager.tabular, next_count)
         self._update_blobs_for_prefix(
-            self._storage_manager.vector, next_count, self._add_vector_indexes
+            next_count,
+            self._storage_manager.tabular_prefix,
+        )
+        self._update_blobs_for_prefix(
+            next_count,
+            self._storage_manager.vector_prefix,
+            self._add_vector_indexes,
         )
         self._table_counts = next_count
 

{corvic_engine-0.3.0rc47 → corvic_engine-0.3.0rc49}/python/corvic/table/table.py

@@ -23,7 +23,7 @@ import structlog
 from google.protobuf import struct_pb2
 from typing_extensions import Self
 
-from corvic import op_graph
+from corvic import op_graph, orm
 from corvic.op_graph import Encoder, Schema
 from corvic.result import (
     InternalError,
@@ -341,17 +341,19 @@ class Table:
         if len(null_columns) > 0:
             _logger.warning("dropped null columns", columns=null_columns)
 
-        match client.storage_manager.tabular.blob_name_from_url(url).and_then(
-            lambda t: op_graph.from_staging(
-                blob_names=[t],
-                arrow_schema=schema.to_arrow(),
-                feature_types=[field.ftype for field in schema],
-                expected_rows=num_rows,
-            )
+        blob_name = client.storage_manager.blob_from_url(blob.url).name.removeprefix(
+            client.storage_manager.tabular_prefix + "/"
+        )
+
+        match op_graph.from_staging(
+            blob_names=[blob_name],
+            arrow_schema=schema.to_arrow(),
+            feature_types=[field.ftype for field in schema],
+            expected_rows=num_rows,
         ):
             case Ok(op):
                 pass
-            case DataMisplacedError() | InvalidArgumentError() as error:
+            case InvalidArgumentError() as error:
                 return error
         if len(null_columns) > 0:
             match op.select_columns(kept_columns):
@@ -366,7 +368,7 @@ class Table:
         return self.op_graph.to_bytes()
 
     def to_polars(
-        self, *, flatten_single_field: bool = False
+        self, room_id: orm.RoomID, *, flatten_single_field: bool = False
     ) -> (
         Ok[Iterable[pl.DataFrame]]
         | InvalidArgumentError
@@ -374,7 +376,7 @@ class Table:
         | ResourceExhaustedError
     ):
         """Stream over the view as a series of Polars DataFrames."""
-        match self.to_batches():
+        match self.to_batches(room_id=room_id):
             case Ok(batch_reader):
                 pass
             case err:
@@ -438,7 +440,7 @@ class Table:
         )
 
     def to_batches(
-        self,
+        self, room_id: orm.RoomID
     ) -> (
         Ok[pa.RecordBatchReader]
         | InvalidArgumentError
@@ -450,7 +452,9 @@ class Table:
             tables_to_compute=[
                 TableComputeContext(
                     self.op_graph,
-                    output_url_prefix=self.client.storage_manager.space_run.make_anonymous_table_url(),
+                    output_url_prefix=self.client.storage_manager.make_tabular_blob(
+                        room_id=room_id, suffix="anonymous_tables"
+                    ).url,
                 )
             ],
         )

{corvic_engine-0.3.0rc47 → corvic_engine-0.3.0rc49}/python/corvic_generated/orm/v1/agent_pb2.py

@@ -14,7 +14,7 @@ _sym_db = _symbol_database.Default()
 
 
 
-DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x19\x63orvic/orm/v1/agent.proto\x12\rcorvic.orm.v1\"\xf3\x01\n\x12\x45xecutorParameters\x12L\n\x15\x63ompletion_model_type\x18\x01 \x01(\x0e\x32\x18.corvic.orm.v1.ModelTypeR\x13\x63ompletionModelType\x12\x42\n\x1d\x63ompletion_system_instruction\x18\x02 \x01(\tR\x1b\x63ompletionSystemInstruction\x12\x33\n\x13\x63ompletion_model_id\x18\x03 \x01(\tH\x00R\x11\x63ompletionModelId\x88\x01\x01\x42\x16\n\x14_completion_model_id\"V\n\x10SpaceInstruction\x12 \n\x0binstruction\x18\x01 \x01(\tR\x0binstruction\x12 \n\x0cspace_run_id\x18\x02 \x01(\tR\nspaceRunId\"\x89\x02\n\x16OrchestratorParameters\x12\x1b\n\tspace_ids\x18\x01 \x03(\tR\x08spaceIds\x12k\n\x12space_instructions\x18\x02 \x03(\x0b\x32<.corvic.orm.v1.OrchestratorParameters.SpaceInstructionsEntryR\x11spaceInstructions\x1a\x65\n\x16SpaceInstructionsEntry\x12\x10\n\x03key\x18\x01 \x01(\tR\x03key\x12\x35\n\x05value\x18\x02 \x01(\x0b\x32\x1f.corvic.orm.v1.SpaceInstructionR\x05value:\x02\x38\x01\"\xc5\x01\n\x0f\x41gentParameters\x12R\n\x13\x65xecutor_parameters\x18\x01 \x01(\x0b\x32!.corvic.orm.v1.ExecutorParametersR\x12\x65xecutorParameters\x12^\n\x17orchestrator_parameters\x18\x02 \x01(\x0b\x32%.corvic.orm.v1.OrchestratorParametersR\x16orchestratorParameters*\xe0\x04\n\tModelType\x12\x1a\n\x16MODEL_TYPE_UNSPECIFIED\x10\x00\x12\x1d\n\x19MODEL_TYPE_GEMINI_1_5_PRO\x10\x01\x12%\n!MODEL_TYPE_GEMINI_1_5_PRO_PREVIEW\x10\x02\x12\x1f\n\x1bMODEL_TYPE_GEMINI_1_5_FLASH\x10\x03\x12\'\n#MODEL_TYPE_GEMINI_1_5_FLASH_PREVIEW\x10\x04\x12\'\n#MODEL_TYPE_GEMINI_2_0_FLASH_PREVIEW\x10\r\x12\x30\n,MODEL_TYPE_GEMINI_2_0_FLASH_THINKING_PREVIEW\x10\x10\x12\x1b\n\x17MODEL_TYPE_GPT_4_O_MINI\x10\x05\x12\x16\n\x12MODEL_TYPE_GPT_4_O\x10\x06\x12\x1c\n\x18MODEL_TYPE_GPT_3_5_TURBO\x10\x07\x12\x12\n\x0eMODEL_TYPE_O_1\x10\x0e\x12\x17\n\x13MODEL_TYPE_O_1_MINI\x10\x0f\x12\x17\n\x13MODEL_TYPE_O_3_MINI\x10\x12\x12 \n\x1cMODEL_TYPE_CLAUDE_3_5_SONNET\x10\x08\x12\x1d\n\x19MODEL_TYPE_CLAUDE_3_HAIKU\x10\t\x12\x1c\n\x18MODEL_TYPE_MISTRAL_LARGE\x10\n\x12\x1b\n\x17MODEL_TYPE_MISTRAL_NEMO\x10\x0b\x12 \n\x1cMODEL_TYPE_MISTRAL_CODESTRAL\x10\x0c\x12\x15\n\x11MODEL_TYPE_CUSTOM\x10\x11\x62\x06proto3')
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x19\x63orvic/orm/v1/agent.proto\x12\rcorvic.orm.v1\"\xf3\x01\n\x12\x45xecutorParameters\x12L\n\x15\x63ompletion_model_type\x18\x01 \x01(\x0e\x32\x18.corvic.orm.v1.ModelTypeR\x13\x63ompletionModelType\x12\x42\n\x1d\x63ompletion_system_instruction\x18\x02 \x01(\tR\x1b\x63ompletionSystemInstruction\x12\x33\n\x13\x63ompletion_model_id\x18\x03 \x01(\tH\x00R\x11\x63ompletionModelId\x88\x01\x01\x42\x16\n\x14_completion_model_id\"V\n\x10SpaceInstruction\x12 \n\x0binstruction\x18\x01 \x01(\tR\x0binstruction\x12 \n\x0cspace_run_id\x18\x02 \x01(\tR\nspaceRunId\"\x89\x02\n\x16OrchestratorParameters\x12\x1b\n\tspace_ids\x18\x01 \x03(\tR\x08spaceIds\x12k\n\x12space_instructions\x18\x02 \x03(\x0b\x32<.corvic.orm.v1.OrchestratorParameters.SpaceInstructionsEntryR\x11spaceInstructions\x1a\x65\n\x16SpaceInstructionsEntry\x12\x10\n\x03key\x18\x01 \x01(\tR\x03key\x12\x35\n\x05value\x18\x02 \x01(\x0b\x32\x1f.corvic.orm.v1.SpaceInstructionR\x05value:\x02\x38\x01\"\xc5\x01\n\x0f\x41gentParameters\x12R\n\x13\x65xecutor_parameters\x18\x01 \x01(\x0b\x32!.corvic.orm.v1.ExecutorParametersR\x12\x65xecutorParameters\x12^\n\x17orchestrator_parameters\x18\x02 \x01(\x0b\x32%.corvic.orm.v1.OrchestratorParametersR\x16orchestratorParameters*\x81\x05\n\tModelType\x12\x1a\n\x16MODEL_TYPE_UNSPECIFIED\x10\x00\x12\x1d\n\x19MODEL_TYPE_GEMINI_1_5_PRO\x10\x01\x12%\n!MODEL_TYPE_GEMINI_1_5_PRO_PREVIEW\x10\x02\x12\x1f\n\x1bMODEL_TYPE_GEMINI_1_5_FLASH\x10\x03\x12\'\n#MODEL_TYPE_GEMINI_1_5_FLASH_PREVIEW\x10\x04\x12\'\n#MODEL_TYPE_GEMINI_2_0_FLASH_PREVIEW\x10\r\x12\x30\n,MODEL_TYPE_GEMINI_2_0_FLASH_THINKING_PREVIEW\x10\x10\x12\x1f\n\x1bMODEL_TYPE_GEMINI_2_0_FLASH\x10\x13\x12\x1b\n\x17MODEL_TYPE_GPT_4_O_MINI\x10\x05\x12\x16\n\x12MODEL_TYPE_GPT_4_O\x10\x06\x12\x1c\n\x18MODEL_TYPE_GPT_3_5_TURBO\x10\x07\x12\x12\n\x0eMODEL_TYPE_O_1\x10\x0e\x12\x17\n\x13MODEL_TYPE_O_1_MINI\x10\x0f\x12\x17\n\x13MODEL_TYPE_O_3_MINI\x10\x12\x12 \n\x1cMODEL_TYPE_CLAUDE_3_5_SONNET\x10\x08\x12\x1d\n\x19MODEL_TYPE_CLAUDE_3_HAIKU\x10\t\x12\x1c\n\x18MODEL_TYPE_MISTRAL_LARGE\x10\n\x12\x1b\n\x17MODEL_TYPE_MISTRAL_NEMO\x10\x0b\x12 \n\x1cMODEL_TYPE_MISTRAL_CODESTRAL\x10\x0c\x12\x15\n\x11MODEL_TYPE_CUSTOM\x10\x11\x62\x06proto3')
 
 _globals = globals()
 _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
@@ -24,7 +24,7 @@ if _descriptor._USE_C_DESCRIPTORS == False:
   _globals['_ORCHESTRATORPARAMETERS_SPACEINSTRUCTIONSENTRY']._options = None
   _globals['_ORCHESTRATORPARAMETERS_SPACEINSTRUCTIONSENTRY']._serialized_options = b'8\001'
   _globals['_MODELTYPE']._serialized_start=847
-  _globals['_MODELTYPE']._serialized_end=1455
+  _globals['_MODELTYPE']._serialized_end=1488
   _globals['_EXECUTORPARAMETERS']._serialized_start=45
   _globals['_EXECUTORPARAMETERS']._serialized_end=288
   _globals['_SPACEINSTRUCTION']._serialized_start=290

{corvic_engine-0.3.0rc47 → corvic_engine-0.3.0rc49}/python/corvic_generated/orm/v1/agent_pb2.pyi

@@ -15,6 +15,7 @@ class ModelType(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
     MODEL_TYPE_GEMINI_1_5_FLASH_PREVIEW: _ClassVar[ModelType]
     MODEL_TYPE_GEMINI_2_0_FLASH_PREVIEW: _ClassVar[ModelType]
     MODEL_TYPE_GEMINI_2_0_FLASH_THINKING_PREVIEW: _ClassVar[ModelType]
+    MODEL_TYPE_GEMINI_2_0_FLASH: _ClassVar[ModelType]
     MODEL_TYPE_GPT_4_O_MINI: _ClassVar[ModelType]
     MODEL_TYPE_GPT_4_O: _ClassVar[ModelType]
     MODEL_TYPE_GPT_3_5_TURBO: _ClassVar[ModelType]
@@ -34,6 +35,7 @@ MODEL_TYPE_GEMINI_1_5_FLASH: ModelType
 MODEL_TYPE_GEMINI_1_5_FLASH_PREVIEW: ModelType
 MODEL_TYPE_GEMINI_2_0_FLASH_PREVIEW: ModelType
 MODEL_TYPE_GEMINI_2_0_FLASH_THINKING_PREVIEW: ModelType
+MODEL_TYPE_GEMINI_2_0_FLASH: ModelType
 MODEL_TYPE_GPT_4_O_MINI: ModelType
 MODEL_TYPE_GPT_4_O: ModelType
 MODEL_TYPE_GPT_3_5_TURBO: ModelType