From 314c951641cef4e89dbdac9a1c740560bb7f6702 Mon Sep 17 00:00:00 2001
From: Matthew Upson <matt@mantisnlp.com>
Date: Thu, 20 Jul 2023 12:32:34 +0200
Subject: [PATCH] fix: Breaking changes introduced in pydantic 2.0.3

---
 hugie/config.py      | 4 ++--
 hugie/endpoint.py    | 4 ++--
 hugie/models.py      | 6 ++++--
 setup.cfg            | 1 +
 tests/test_config.py | 2 +-
 5 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/hugie/config.py b/hugie/config.py
index 6d0a296..30db66a 100644
--- a/hugie/config.py
+++ b/hugie/config.py
@@ -100,13 +100,13 @@ def modify(
 
     if overwrite:
         try:
-            srsly.write_json(path, config.dict())
+            srsly.write_json(path, config.model_dump())
             typer.secho(f"Updated config at {path}", fg=typer.colors.GREEN)
         except Exception:
             typer.secho(f"Failed to update config at {path}", fg=typer.colors.RED)
 
     else:
-        typer.secho(srsly.json_dumps(config.dict()), fg=typer.colors.YELLOW)
+        typer.secho(srsly.json_dumps(config.model_dump()), fg=typer.colors.YELLOW)
 
 
 if __name__ == "__main__":
diff --git a/hugie/endpoint.py b/hugie/endpoint.py
index 1994d2d..4765e28 100644
--- a/hugie/endpoint.py
+++ b/hugie/endpoint.py
@@ -88,7 +88,7 @@ def create(
         data (str): Path to JSON data to create the endpoint
     """
 
-    data = InferenceEndpointConfig.from_json(data).dict()
+    data = InferenceEndpointConfig.from_json(data).model_dump()
 
     try:
         response = requests.post(settings.endpoint_url, headers=headers, json=data)
@@ -126,7 +126,7 @@ def update(
     """
     Update an endpoint
     """
-    data = InferenceEndpointConfig.from_json(data).dict()
+    data = InferenceEndpointConfig.from_json(data).model_dump()
 
     try:
         response = requests.put(
diff --git a/hugie/models.py b/hugie/models.py
index eb41aaf..26c3f14 100644
--- a/hugie/models.py
+++ b/hugie/models.py
@@ -1,4 +1,6 @@
-from pydantic import BaseModel, BaseSettings
+from pydantic import BaseModel
+from pydantic_settings import BaseSettings
+from typing import Optional
 
 from hugie.utils import load_json
 
@@ -35,7 +37,7 @@ class InferenceEndpointConfig(BaseSettings):
     Config for the inference endpoint
     """
 
-    accountId: str = None
+    accountId: Optional[str] = None
     type: str = None
     compute: ComputeModel = ComputeModel()
     model: ModelModel = ModelModel()
diff --git a/setup.cfg b/setup.cfg
index 33868cb..bebd494 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -20,6 +20,7 @@ install_requires =
     requests
     pydantic
     srsly
+    pydantic-settings
 python_requires = >= 3.6
 
 
diff --git a/tests/test_config.py b/tests/test_config.py
index 1b335a1..f42e470 100644
--- a/tests/test_config.py
+++ b/tests/test_config.py
@@ -40,7 +40,7 @@ def data_path(tmp_path):
 
 def test_inference_endpoint_config_serialization(data_path):
     config = InferenceEndpointConfig.from_json(data_path)
-    config = config.dict()
+    config = config.model_dump()
     assert isinstance(config, dict)
     assert config["compute"]["instanceSize"] == "small"
     assert config["model"]["framework"] == "custom"