Removed JSONType for Pydantic's equivalent JsonValue (#1022)

Author: jamesbraza

paperqa/agents/models.py

@@ -5,28 +5,18 @@
 import time
 from contextlib import asynccontextmanager
 from enum import StrEnum
-from typing import Any, ClassVar, Protocol, TypeAlias, cast
+from typing import Any, ClassVar, Protocol, cast
 from uuid import UUID, uuid4
 
 from aviary.core import Message
 from lmi import LiteLLMModel, LLMModel
-from pydantic import (
-    BaseModel,
-    ConfigDict,
-    Field,
-    ValidationInfo,
-    field_validator,
-)
+from pydantic import BaseModel, ConfigDict, Field, ValidationInfo, field_validator
 
 from paperqa.types import PQASession
 from paperqa.version import __version__
 
 logger = logging.getLogger(__name__)
 
-JSONType: TypeAlias = (
-    dict[str, "JSONType"] | list["JSONType"] | str | int | float | bool | None
-)
-
 
 class SupportsPickle(Protocol):
     """Type protocol for typing any object that supports pickling."""

paperqa/agents/search.py

@@ -17,7 +17,7 @@
 from typing import TYPE_CHECKING, Any, ClassVar
 
 import anyio
-from pydantic import BaseModel
+from pydantic import BaseModel, JsonValue
 from rich.progress import (
     BarColumn,
     MofNCompleteColumn,
@@ -47,7 +47,7 @@
 from paperqa.types import VAR_MATCH_LOOKUP, DocDetails
 from paperqa.utils import ImpossibleParsingError, clean_possessives, hexdigest
 
-from .models import JSONType, SupportsPickle
+from .models import SupportsPickle
 
 if TYPE_CHECKING:
     from tantivy import IndexWriter
@@ -86,7 +86,7 @@ def write_to_string(self, data: BaseModel | SupportsPickle) -> bytes:
 
     def read_from_string(
         self, data: str | bytes
-    ) -> BaseModel | SupportsPickle | JSONType:
+    ) -> BaseModel | SupportsPickle | JsonValue:
         if self == SearchDocumentStorage.JSON_MODEL_DUMP:
             return json.loads(data)
         if self == SearchDocumentStorage.PICKLE_COMPRESSED: