Modernise

great_ai/context.py

@@ -4,7 +4,7 @@ from logging import DEBUG, Logger
 from pathlib import Path
 from typing import Any, Dict, Optional, Type, Union, cast
 
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict
 
 from great_ai import __version__
 
@@ -34,8 +34,7 @@ class Context(BaseModel):
     dashboard_table_size: int
     route_config: RouteConfig
 
-    class Config:
-        arbitrary_types_allowed = True
+    model_config = ConfigDict(arbitrary_types_allowed=True)
 
     def to_flat_dict(self) -> Dict[str, Any]:
         return {
@@ -137,9 +136,11 @@ def configure(
         ),
         is_production=is_production,
         logger=logger,
-        should_log_exception_stack=not is_production
-        if should_log_exception_stack is None
-        else should_log_exception_stack,
+        should_log_exception_stack=(
+            not is_production
+            if should_log_exception_stack is None
+            else should_log_exception_stack
+        ),
         prediction_cache_size=prediction_cache_size,
         dashboard_table_size=dashboard_table_size,
         route_config=route_config,

great_ai/deploy/great_ai.py

@@ -16,7 +16,7 @@ from typing import (
 )
 
 from fastapi import FastAPI
-from tqdm.cli import tqdm
+from tqdm import tqdm
 from typing_extensions import Literal  # <= Python 3.7
 
 from ..constants import DASHBOARD_PATH
@@ -103,15 +103,13 @@ class GreatAI(Generic[T, V]):
         # Overloaded function signatures 1 and 2 overlap with incompatible return types
         # https://github.com/python/mypy/issues/12759
         func: Callable[..., Awaitable[V]],
-    ) -> "GreatAI[Awaitable[Trace[V]], V]":
-        ...
+    ) -> "GreatAI[Awaitable[Trace[V]], V]": ...
 
     @overload
     @staticmethod
     def create(
         func: Callable[..., V],
-    ) -> "GreatAI[Trace[V], V]":
-        ...
+    ) -> "GreatAI[Trace[V], V]": ...
 
     @staticmethod
     def create(
@@ -146,7 +144,7 @@ class GreatAI(Generic[T, V]):
             >>> my_function('3').output
             Traceback (most recent call last):
                 ...
-            TypeError: type of a must be int; got str instead
+            TypeError: argument a is not of the expected type: ...
 
         Args:
             func: The prediction function that needs to be decorated.
@@ -170,8 +168,7 @@ class GreatAI(Generic[T, V]):
         concurrency: Optional[int] = None,
         unpack_arguments: Literal[True],
         do_not_persist_traces: bool = ...,
-    ) -> List[Trace[V]]:
-        ...
+    ) -> List[Trace[V]]: ...
 
     @overload
     def process_batch(
@@ -181,8 +178,7 @@ class GreatAI(Generic[T, V]):
         concurrency: Optional[int] = None,
         unpack_arguments: Literal[False] = ...,
         do_not_persist_traces: bool = ...,
-    ) -> List[Trace[V]]:
-        ...
+    ) -> List[Trace[V]]: ...
 
     def process_batch(
         self,
@@ -237,9 +233,11 @@ class GreatAI(Generic[T, V]):
         return list(
             tqdm(
                 parallel_map(
-                    inner_async
-                    if get_function_metadata_store(self).is_asynchronous
-                    else inner,
+                    (
+                        inner_async
+                        if get_function_metadata_store(self).is_asynchronous
+                        else inner
+                    ),
                     batch,
                     concurrency=concurrency,
                 ),
@@ -249,7 +247,7 @@ class GreatAI(Generic[T, V]):
 
     @staticmethod
     def _get_cached_traced_function(
-        func: Callable[..., Union[V, Awaitable[V]]]
+        func: Callable[..., Union[V, Awaitable[V]]],
     ) -> Callable[..., T]:
         @lru_cache(maxsize=get_context().prediction_cache_size)
         def func_in_tracing_context_sync(

great_ai/deploy/routes/bootstrap_feedback_endpoints.py

@@ -31,7 +31,7 @@ def bootstrap_feedback_endpoints(app: FastAPI) -> None:
         return trace.feedback
 
     @router.delete("/", status_code=status.HTTP_204_NO_CONTENT)
-    def delete_feedback(trace_id: str) -> Any:
+    def delete_feedback(trace_id: str) -> Response:
         trace = get_context().tracing_database.get(trace_id)
         if trace is None:
             raise HTTPException(status_code=status.HTTP_404_NOT_FOUND)

great_ai/deploy/routes/bootstrap_prediction_endpoint.py

@@ -22,10 +22,10 @@ def bootstrap_prediction_endpoint(
         try:
             if inspect.iscoroutinefunction(func):
                 return await cast(Callable[..., Awaitable[Trace]], func)(
-                    **cast(BaseModel, input_value).dict()
+                    **cast(BaseModel, input_value).model_dump()
                 )
             return cast(Callable[..., Trace], func)(
-                **cast(BaseModel, input_value).dict()
+                **cast(BaseModel, input_value).model_dump()
             )
         except Exception as e:
             raise HTTPException(

great_ai/deploy/routes/dashboard/create_dash_app.py

@@ -26,9 +26,9 @@ def create_dash_app(function_name: str, version: str, function_docs: str) -> Fla
     app = Dash(
         function_name,
         requests_pathname_prefix=DASHBOARD_PATH + "/",
-        server=Flask(__name__),
+        server=Flask(__name__),  # type: ignore[arg-type]
         title=function_name,
-        update_title=None,
+        update_title=None,  # type: ignore[arg-type]
         external_stylesheets=[
             "/assets/index.css",
         ],
@@ -64,9 +64,9 @@ def create_dash_app(function_name: str, version: str, function_docs: str) -> Fla
     interval = dcc.Interval(
         interval=2 * 1000,  # in milliseconds
         n_intervals=0,
-        max_intervals=1
-        if get_context().is_production
-        else -1,  # will be incremented in production upon each successful request
+        max_intervals=(
+            1 if get_context().is_production else -1
+        ),  # will be incremented in production upon each successful request
     )
 
     app.layout = html.Main(
@@ -159,7 +159,7 @@ def create_dash_app(function_name: str, version: str, function_docs: str) -> Fla
             take=page_size,
             conjunctive_filters=non_null_conjunctive_filters,
             conjunctive_tags=[ONLINE_TAG_NAME],
-            sort_by=[SortBy.parse_obj(s) for s in sort_by],
+            sort_by=[SortBy.model_validate(s) for s in sort_by],
         )
 
         if non_null_conjunctive_filters:
@@ -244,7 +244,7 @@ def update_charts(
     fig.update_layout(
         margin=dict(l=0, r=0, b=0, t=0, pad=0),
     )
-    execution_time_histogram.figure = fig
+    execution_time_histogram.figure = fig  # type: ignore[attr-defined]
 
     parallel_coords_fig = go.Figure(
         go.Parcoords(

great_ai/deploy/routes/dashboard/get_description.py

@@ -16,8 +16,7 @@ def get_description(
                 style={"color": accent_color},
             ),
             dcc.Markdown(
-                strip_lines(
-                    f"""
+                strip_lines(f"""
                     > View the live data of your deployment here.
 
                     ## Using the API
@@ -27,8 +26,7 @@ def get_description(
                     ## Details
 
                     {function_docs}
-                """
-                ),
+                """),
                 className="description",
             ),
         ]

great_ai/deploy/routes/dashboard/get_traces_table.py

@@ -1,10 +1,12 @@
+from typing import Any
+
 from dash import dash_table
 
 from ....context import get_context
 
 
-def get_traces_table() -> dash_table.DataTable:
-    return dash_table.DataTable(
+def get_traces_table() -> Any:
+    return dash_table.DataTable(  # type: ignore[attr-defined]
         page_current=0,
         page_size=get_context().dashboard_table_size,
         page_action="custom",

great_ai/helper/freeze_arguments.py

@@ -1,12 +1,24 @@
 import inspect
-from functools import wraps
-from typing import Any, Callable, Mapping, Sequence, Set, TypeVar, Union, cast
+from functools import lru_cache, wraps
+from typing import Any, Callable, Mapping, Sequence, Set, Type, TypeVar, Union, cast
 
 from pydantic import BaseModel
 
 F = TypeVar("F", bound=Callable)
 
 
+@lru_cache(maxsize=None)
+def _hashable_model_type(model_type: Type[BaseModel]) -> Type[BaseModel]:
+    # The subclass is cached per model type: Pydantic v2's __eq__ requires both
+    # operands to share a type, so creating a fresh subclass on every freeze()
+    # call would make even identical models compare unequal.
+    class HashableValue(model_type):  # type: ignore
+        def __hash__(self) -> int:
+            return hash(frozenset((k, freeze(v)) for k, v in self.model_dump().items()))
+
+    return HashableValue
+
+
 class FrozenDict(dict):
     def __hash__(self) -> int:  # type: ignore
         return hash(frozenset((k, freeze(v)) for k, v in self.items()))
@@ -55,11 +67,6 @@ def freeze(value: Union[Sequence[Any], Mapping[str, Any], Set[Any], BaseModel])
         return FrozenSet(value)
 
     if isinstance(value, BaseModel):
-
-        class HashableValue(type(value)):  # type: ignore
-            def __hash__(self) -> int:
-                return hash(frozenset((k, freeze(v)) for k, v in self.dict().items()))
-
-        return HashableValue(**value.dict())
+        return _hashable_model_type(type(value))(**value.model_dump())
 
     return value

great_ai/large_file/helper/cached_property.py

@@ -28,7 +28,9 @@ class cached_property:
             )
         try:
             cache = instance.__dict__
-        except AttributeError:  # not all objects have __dict__ (e.g. class defines slots)
+        except (
+            AttributeError
+        ):  # not all objects have __dict__ (e.g. class defines slots)
             msg = (
                 f"No '__dict__' attribute on {type(instance).__name__!r} "
                 f"instance to cache {self.attrname!r} property."

great_ai/large_file/large_file/large_file_mongo.py

@@ -2,8 +2,9 @@ import re
 from pathlib import Path
 from typing import Any, List
 
-from gridfs import DEFAULT_CHUNK_SIZE, Database, GridFSBucket
+from gridfs import DEFAULT_CHUNK_SIZE, GridFSBucket
 from pymongo import MongoClient
+from pymongo.database import Database
 
 from ...utilities import get_logger
 from ..helper import DownloadProgressBar, UploadProgressBar, cached_property
@@ -60,7 +61,6 @@ class LargeFileMongo(LargeFileBase):
                 ),
                 version=int(f.name.split(MONGO_NAME_VERSION_SEPARATOR)[-1]),
                 remote_path=(f._id, f.length),
-                origin="mongodb",
             )
             for f in self._client.find(
                 {

great_ai/large_file/large_file/large_file_s3.py

@@ -100,9 +100,13 @@ class LargeFileS3(LargeFileBase):
             Bucket=self.bucket_name,
             Key=remote_path,
             Filename=str(local_path),
-            Callback=None
-            if hide_progress
-            else DownloadProgressBar(name=str(remote_path), size=size, logger=logger),
+            Callback=(
+                None
+                if hide_progress
+                else DownloadProgressBar(
+                    name=str(remote_path), size=size, logger=logger
+                )
+            ),
         )
 
     def _upload(self, local_path: Path, hide_progress: bool) -> None:
@@ -113,9 +117,11 @@ class LargeFileS3(LargeFileBase):
             Filename=str(local_path),
             Bucket=self.bucket_name,
             Key=key,
-            Callback=None
-            if hide_progress
-            else UploadProgressBar(path=local_path, logger=logger),
+            Callback=(
+                None
+                if hide_progress
+                else UploadProgressBar(path=local_path, logger=logger)
+            ),
         )
 
     def _delete_old_remote_versions(self) -> None:

great_ai/parameters/parameter.py

@@ -1,7 +1,7 @@
 from functools import wraps
 from typing import Any, Callable, Dict, TypeVar, cast
 
-from typeguard import check_type
+from typeguard import TypeCheckError, check_type
 
 from ..errors import ArgumentValidationError
 from ..helper import get_arguments, get_function_metadata_store
@@ -32,7 +32,7 @@ def parameter(
         >>> my_function('3')
         Traceback (most recent call last):
             ...
-        TypeError: type of a must be int; got str instead
+        TypeError: argument a is not of the expected type: ...
 
         >>> @parameter('positive_a', validate=lambda v: v > 0)
         ... def my_function(positive_a: int):
@@ -65,13 +65,16 @@ def parameter(
             expected_type = func.__annotations__.get(parameter_name)
 
             if expected_type is not None:
-                check_type(parameter_name, argument, expected_type)
+                try:
+                    check_type(argument, expected_type)
+                except TypeCheckError as error:
+                    raise TypeError(
+                        f"argument {parameter_name} is not of the expected type: {error}"
+                    ) from error
 
             if not validate(argument):
                 raise ArgumentValidationError(
-                    f"""Argument {parameter_name} in {
-                        func.__name__
-                    } did not pass validation"""
+                    f"Argument {parameter_name} in {func.__name__} did not pass validation"
                 )
 
             context = TracingContext.get_current_tracing_context()

great_ai/persistence/mongodb_driver.py

@@ -90,7 +90,7 @@ class MongoDbDriver(TracingDatabaseDriver):
             value = client[self.mongo_database].traces.find_one(id)
 
         if value:
-            value = Trace.parse_obj(value)
+            value = Trace.model_validate(value)
 
         return value
 
@@ -150,7 +150,7 @@ class MongoDbDriver(TracingDatabaseDriver):
             ]
 
             with client[self.mongo_database].traces.find(**query) as cursor:
-                documents = [Trace[Any].parse_obj(t) for t in cursor]
+                documents = [Trace[Any].model_validate(t) for t in cursor]
         return documents, count
 
     def update(self, id: str, new_version: Trace) -> None:

great_ai/persistence/parallel_tinydb_driver.py

@@ -29,16 +29,16 @@ class ParallelTinyDbDriver(TracingDatabaseDriver):
     path_to_db = Path(DEFAULT_TRACING_DB_FILENAME)
 
     def save(self, trace: Trace) -> str:
-        return self._safe_execute(lambda db: db.insert(trace.dict()))
+        return self._safe_execute(lambda db: db.insert(trace.model_dump()))
 
     def save_batch(self, documents: List[Trace]) -> List[str]:
-        traces = [d.dict() for d in documents]
+        traces = [d.model_dump() for d in documents]
         return self._safe_execute(lambda db: db.insert_multiple(traces))
 
     def get(self, id: str) -> Optional[Trace]:
         value = self._safe_execute(lambda db: db.get(lambda d: d["trace_id"] == id))
         if value:
-            value = Trace.parse_obj(value)
+            value = Trace.model_validate(value)
         return value
 
     def query(
@@ -51,7 +51,7 @@ class ParallelTinyDbDriver(TracingDatabaseDriver):
         since: Optional[datetime] = None,
         until: Optional[datetime] = None,
         has_feedback: Optional[bool] = None,
-        sort_by: Sequence[SortBy] = []
+        sort_by: Sequence[SortBy] = [],
     ) -> Tuple[List[Trace], int]:
         def does_match(d: Dict[str, Any]) -> bool:
             return (
@@ -67,7 +67,11 @@ class ParallelTinyDbDriver(TracingDatabaseDriver):
         if not documents:
             return [], 0
 
-        df = pd.DataFrame([Trace.parse_obj(d).to_flat_dict() for d in documents])
+        traces: List[Trace] = [Trace.model_validate(d) for d in documents]
+        # The DataFrame keeps its default 0..n-1 index, so filtering/sorting below
+        # preserve the row labels that index back into `traces`. This avoids
+        # reconstructing Traces from the lossy flattened rows.
+        df = pd.DataFrame([t.to_flat_dict() for t in traces])
 
         for f in conjunctive_filters:
             operator = f.operator.lower()
@@ -92,11 +96,13 @@ class ParallelTinyDbDriver(TracingDatabaseDriver):
 
         count = len(df)
         result = df.iloc[skip:] if take is None else df.iloc[skip : skip + take]
-        return [Trace.parse_obj(trace) for _, trace in result.iterrows()], count
+        return [traces[i] for i in result.index], count
 
     def update(self, id: str, new_version: Trace) -> None:
         self._safe_execute(
-            lambda db: db.update(new_version.dict(), lambda d: d["trace_id"] == id)
+            lambda db: db.update(
+                new_version.model_dump(), lambda d: d["trace_id"] == id
+            )
         )
 
     def delete(self, id: str) -> None:

great_ai/persistence/tracing_database_driver.py

@@ -54,7 +54,7 @@ class TracingDatabaseDriver(ABC):
         until: Optional[datetime] = None,
         since: Optional[datetime] = None,
         has_feedback: Optional[bool] = None,
-        sort_by: Sequence[SortBy] = []
+        sort_by: Sequence[SortBy] = [],
     ) -> Tuple[List[Trace], int]:
         pass
 

great_ai/remote/call_remote_great_ai_async.py

@@ -64,7 +64,7 @@ async def call_remote_great_ai_async(
         )
     try:
         if model_class is not None:
-            trace["output"] = model_class.parse_obj(trace["output"])
-        return Trace.parse_obj(trace)
+            trace["output"] = model_class.model_validate(trace["output"])
+        return Trace.model_validate(trace)
     except Exception:
         raise RemoteCallError("Could not parse Trace")

great_ai/tracing/add_ground_truth.py

@@ -23,7 +23,7 @@ def add_ground_truth(
     tags: Union[List[str], str] = [],
     train_split_ratio: float = 1,
     test_split_ratio: float = 0,
-    validation_split_ratio: float = 0
+    validation_split_ratio: float = 0,
 ) -> None:
     """Add training data (with optional train-test splitting).
 

great_ai/tracing/query_ground_truth.py

@@ -10,7 +10,7 @@ def query_ground_truth(
     *,
     since: Optional[datetime] = None,
     until: Optional[datetime] = None,
-    return_max_count: Optional[int] = None
+    return_max_count: Optional[int] = None,
 ) -> List[Trace]:
     """Return training samples.
 

great_ai/tracing/tracing_context.py

@@ -51,15 +51,19 @@ class TracingContext(Generic[T]):
                     logged_values=self._values,
                     models=self._models,
                     output=output,
-                    exception=None
-                    if exception is None
-                    else f"{type(exception).__name__}: {exception}",
+                    exception=(
+                        None
+                        if exception is None
+                        else f"{type(exception).__name__}: {exception}"
+                    ),
                     tags=[
                         self._name,
                         ONLINE_TAG_NAME,
-                        PRODUCTION_TAG_NAME
-                        if get_context().is_production
-                        else DEVELOPMENT_TAG_NAME,
+                        (
+                            PRODUCTION_TAG_NAME
+                            if get_context().is_production
+                            else DEVELOPMENT_TAG_NAME
+                        ),
                     ],
                 ),
             )

great_ai/utilities/config_file/config_file.py

@@ -94,11 +94,9 @@ class ConfigFile(Mapping[str, str]):
             try:
                 value = next(v for v in values if v)
             except StopIteration:
-                raise ParseError(
-                    f"""Cannot parse config file ({
+                raise ParseError(f"""Cannot parse config file ({
                         self._path.absolute()
-                        }), error at key `{key}`"""
-                )
+                        }), error at key `{key}`""")
 
             already_exists = key in self._key_values
             if already_exists and not value.startswith(