feat: autogen types from openapi scope3ai yaml

scope3data · Dec 17, 2024 · 43393da · 43393da
1 parent bc36d55
commit 43393da
Show file tree

Hide file tree

Showing 10 changed files with 905 additions and 369 deletions.
diff --git a/README.md b/README.md
@@ -137,3 +137,15 @@ $ export UV_ENV_FILE=.env
 $ uv sync --all-extras --all-groups
 $ uv run python -m examples.openai-sync-chat
 ```
+
+## Update typesgen.py
+
+```bash
+$ uv run datamodel-codegen \
+    --input tests/api-mocks/aiapi.yaml \
+    --input-file-type openapi \
+    --output scope3ai/api/typesgen.py \
+    --output-model-type pydantic_v2.BaseModel \
+    --use-schema-description \
+    --allow-extra-fields
+```
diff --git a/pyproject.toml b/pyproject.toml
@@ -43,3 +43,6 @@ tests = [
     "pytest>=8.3.4",
     "pytest-docker>=3.1.1",
 ]
+openapi = [
+    "datamodel-code-generator>=0.26.4",
+]
diff --git a/scope3ai/api/client.py b/scope3ai/api/client.py
@@ -6,15 +6,15 @@
 
 from .defaults import DEFAULT_API_URL
 from .types import (
-    GpuResponse,
+    GPU,
     ImpactRequest,
-    ImpactRequestRow,
+    ImpactRow,
     ImpactResponse,
-    ModelFamily,
+    Family,
     ModelResponse,
-    NodeCloud,
+    CloudProvider,
     NodeResponse,
-    NodeService,
+    ManagedServiceProvider,
 )
 
 
@@ -53,7 +53,7 @@ def client(self) -> Union[httpx.Client, httpx.AsyncClient]:
 class ClientCommands:
     def model(
         self,
-        family: Optional[ModelFamily] = None,
+        family: Optional[Family] = None,
         with_response: Optional[bool] = True,
     ) -> ModelResponse:
         """
@@ -73,21 +73,21 @@ def model(
     def gpu(
         self,
         with_response: Optional[bool] = True,
-    ) -> GpuResponse:
+    ) -> GPU:
         """
         List GPUs
         """
         return self.execute_request(
             "/gpu",
             method="GET",
-            response_model=GpuResponse,
+            response_model=GPU,
             with_response=with_response,
         )
 
     def node(
         self,
-        service: Optional[NodeService] = None,
-        cloud: Optional[NodeCloud] = None,
+        service: Optional[ManagedServiceProvider] = None,
+        cloud: Optional[CloudProvider] = None,
         with_response: Optional[bool] = True,
     ) -> NodeResponse:
         """
@@ -108,7 +108,7 @@ def node(
 
     def impact(
         self,
-        rows: List[ImpactRequestRow],
+        rows: List[ImpactRow],
         debug: Optional[bool] = None,
         with_response: Optional[bool] = True,
     ) -> ImpactResponse:
@@ -201,17 +201,17 @@ async def execute_request(
 
     funcs_to_test = [
         ["model", {}],
-        ["model", {"family": ModelFamily.CLAUDE}],
+        ["model", {"family": Family.claud}],
         ["gpu", {}],
         ["node", {}],
-        ["node", {"cloud": NodeCloud.AWS}],
-        # ["node", {"service": NodeService.AZURE_ML}],
+        ["node", {"cloud": CloudProvider.aws}],
+        # ["node", {"service": ManagedServiceProvider.azure_ml}],
         ["impact", {"rows": []}],
         [
             "impact",
             {
                 "rows": [
-                    ImpactRequestRow(
+                    ImpactRow(
                         model=Model(id="gpt_4o"),
                         input_tokens=1000,
                         output_tokens=200,