Upgrade to pydantic 2 using bump-pydantic, refs #74

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:

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

setup.py

@@ -37,7 +37,7 @@ setup(
         "openai",
         "click-default-group-wheel",
         "sqlite-utils",
-        "pydantic<2.0",
+        "pydantic>=2.0.0",
         "PyYAML",
         "pluggy",
     ],