refactor: Add types for model downloading

4b3e60f899a5f7f7b8693e4e851db57b7d8046f4 · 2025-06-29 21:59:06 +0100 · Alex Cheema

Files touched

M shared/types/events.py
A shared/types/model.py

Diff

commit 4b3e60f899a5f7f7b8693e4e851db57b7d8046f4
Author: Alex Cheema <alexcheema123@gmail.com>
Date:   Sun Jun 29 21:59:06 2025 +0100

    refactor: Add types for model downloading
---
 shared/types/events.py |  5 +----
 shared/types/model.py  | 47 +++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 48 insertions(+), 4 deletions(-)

diff --git a/shared/types/events.py b/shared/types/events.py
index 6051c0ae..f7ee52a8 100644
--- a/shared/types/events.py
+++ b/shared/types/events.py
@@ -6,10 +6,7 @@ from uuid import UUID
 from pydantic import BaseModel, TypeAdapter, UuidVersion
 
 from shared.types.event_sourcing import Event
-
-_ModelId = Annotated[UUID, UuidVersion(4)]
-ModelId = type("ModelId", (UUID,), {})
-ModelIdParser: TypeAdapter[ModelId] = TypeAdapter(_ModelId)
+from shared.types.model import ModelId
 
 _NodeId = Annotated[UUID, UuidVersion(4)]
 NodeId = type("NodeId", (UUID,), {})
diff --git a/shared/types/model.py b/shared/types/model.py
new file mode 100644
index 00000000..953c333f
--- /dev/null
+++ b/shared/types/model.py
@@ -0,0 +1,47 @@
+from typing import Any, Literal, Annotated, final, TypeVar, Generic
+from pydantic import Field, BaseModel, AnyHttpUrl, TypeAdapter
+from uuid import UUID
+from pydantic.types import UuidVersion
+
+SourceType = Literal["HuggingFace", "GitHub"]
+
+T = TypeVar("T", bound=SourceType)
+
+_ModelId = Annotated[UUID, UuidVersion(4)]
+ModelId = type("ModelId", (UUID,), {})
+ModelIdParser: TypeAdapter[ModelId] = TypeAdapter(_ModelId)
+
+RepoPath = Annotated[str, Field(pattern=r'^[^/]+/[^/]+$')]
+RepoURL = Annotated[str, AnyHttpUrl]
+
+class BaseModelSource(BaseModel, Generic[T]):
+    model_uuid: ModelId
+    source_type: T
+    source_data: Any
+
+@final
+class HuggingFaceModelSourceData(BaseModel):
+    path: RepoPath
+
+@final
+class GitHubModelSourceData(BaseModel):
+    url: AnyHttpUrl
+
+@final
+class HuggingFaceModelSource(BaseModelSource[Literal["HuggingFace"]]):
+    source_type: Literal["HuggingFace"] = "HuggingFace"
+    source_data: HuggingFaceModelSourceData
+
+@final
+class GitHubModelSource(BaseModelSource[Literal["GitHub"]]):
+    source_type: Literal["GitHub"] = "GitHub"
+    source_data: GitHubModelSourceData
+
+RepoType = BaseModelSource[SourceType]
+
+RepoValidatorThing = Annotated[
+    RepoType,
+    Field(discriminator="source_type")
+]
+
+RepoValidator: TypeAdapter[RepoValidatorThing] = TypeAdapter(RepoValidatorThing)

← 784f0ec4 chore: Skip protobuf generation if no .proto files exist · back to Exo · refactor: Use official OpenAI types 0c46adc2 →