From 9ac120fae46496a08f0766c058e4b71c33ef0161 Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Sat, 1 Jul 2023 11:53:54 -0700
Subject: [PATCH] Upgrade to pydantic 2 using bump-pydantic, refs #74

---
 llm/models.py    |  5 ++---
 llm/templates.py | 14 ++++++--------
 setup.py         |  2 +-
 3 files changed, 9 insertions(+), 12 deletions(-)

diff --git a/llm/models.py b/llm/models.py
index 01baae8..f6cae4a 100644
--- a/llm/models.py
+++ b/llm/models.py
@@ -2,7 +2,7 @@ from dataclasses import dataclass
 from typing import Any, Dict, Generator, Optional, Set
 from abc import ABC, abstractmethod
 import os
-from pydantic import BaseModel
+from pydantic import ConfigDict, BaseModel
 
 
 @dataclass
@@ -62,8 +62,7 @@ class Model(ABC):
     can_stream: bool = False
 
     class Options(BaseModel):
-        class Config:
-            extra = "forbid"
+        model_config = ConfigDict(extra="forbid")
 
     def get_key(self):
         if self.needs_key is None:
diff --git a/llm/templates.py b/llm/templates.py
index 4eedbb6..6885fe3 100644
--- a/llm/templates.py
+++ b/llm/templates.py
@@ -1,17 +1,15 @@
-from pydantic import BaseModel
+from pydantic import ConfigDict, BaseModel
 import string
 from typing import Optional
 
 
 class Template(BaseModel):
     name: str
-    prompt: Optional[str]
-    system: Optional[str]
-    model: Optional[str]
-    defaults: Optional[dict]
-
-    class Config:
-        extra = "forbid"
+    prompt: Optional[str] = None
+    system: Optional[str] = None
+    model: Optional[str] = None
+    defaults: Optional[dict] = None
+    model_config = ConfigDict(extra="forbid")
 
     class MissingVariables(Exception):
         pass
diff --git a/setup.py b/setup.py
index f44b792..f8c46c6 100644
--- a/setup.py
+++ b/setup.py
@@ -37,7 +37,7 @@ setup(
         "openai",
         "click-default-group-wheel",
         "sqlite-utils",
-        "pydantic<2.0",
+        "pydantic>=2.0.0",
         "PyYAML",
         "pluggy",
     ],