prediction-market-agent-tooling 0.65.5__py3-none-any.whl → 0.69.17.dev1149__py3-none-any.whl

prediction_market_agent_tooling/tools/caches/db_cache.py

@@ -1,6 +1,8 @@
+import asyncio
 import hashlib
 import inspect
 import json
+from dataclasses import dataclass
 from datetime import timedelta
 from functools import wraps
 from types import UnionType
@@ -12,6 +14,7 @@ from typing import (
     cast,
     get_args,
     get_origin,
+    get_type_hints,
     overload,
 )
 
@@ -25,7 +28,10 @@ from sqlmodel import Field, SQLModel, desc, select
 from prediction_market_agent_tooling.config import APIKeys
 from prediction_market_agent_tooling.loggers import logger
 from prediction_market_agent_tooling.tools.datetime_utc import DatetimeUTC
-from prediction_market_agent_tooling.tools.db.db_manager import DBManager
+from prediction_market_agent_tooling.tools.db.db_manager import (
+    DBManager,
+    EnsureTableManager,
+)
 from prediction_market_agent_tooling.tools.utils import utcnow
 
 DB_CACHE_LOG_PREFIX = "[db-cache]"
@@ -46,6 +52,10 @@ class FunctionCache(SQLModel, table=True):
     created_at: DatetimeUTC = Field(default_factory=utcnow, index=True)
 
 
+# Global instance of the table manager for FunctionCache
+_table_manager = EnsureTableManager([FunctionCache])
+
+
 @overload
 def db_cache(
     func: None = None,
@@ -101,136 +111,228 @@ def db_cache(
 
     api_keys = api_keys if api_keys is not None else APIKeys()
 
-    @wraps(func)
-    def wrapper(*args: Any, **kwargs: Any) -> Any:
-        # If caching is disabled, just call the function and return it
-        if not api_keys.ENABLE_CACHE:
-            return func(*args, **kwargs)
-
-        DBManager(api_keys.sqlalchemy_db_url.get_secret_value()).create_tables(
-            [FunctionCache]
-        )
+    # Check if the decorated function is async
+    if inspect.iscoroutinefunction(func):
 
-        # Convert *args and **kwargs to a single dictionary, where we have names for arguments passed as args as well.
-        signature = inspect.signature(func)
-        bound_arguments = signature.bind(*args, **kwargs)
-        bound_arguments.apply_defaults()
-
-        # Convert any argument that is Pydantic model into classic dictionary, otherwise it won't be json-serializable.
-        args_dict: dict[str, Any] = bound_arguments.arguments
-
-        # Remove `self` or `cls` if present (in case of class' methods)
-        if "self" in args_dict:
-            del args_dict["self"]
-        if "cls" in args_dict:
-            del args_dict["cls"]
-
-        # Remove ignored arguments
-        if ignore_args:
-            for arg in ignore_args:
-                if arg in args_dict:
-                    del args_dict[arg]
-
-        # Remove arguments of ignored types
-        if ignore_arg_types:
-            args_dict = {
-                k: v
-                for k, v in args_dict.items()
-                if not isinstance(v, tuple(ignore_arg_types))
-            }
+        @wraps(func)
+        async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
+            # If caching is disabled, just call the function and return it
+            if not api_keys.ENABLE_CACHE:
+                return await func(*args, **kwargs)
 
-        # Compute a hash of the function arguments used for lookup of cached results
-        arg_string = json.dumps(args_dict, sort_keys=True, default=str)
-        args_hash = hashlib.md5(arg_string.encode()).hexdigest()
+            # Ensure tables are created before accessing cache
+            await _table_manager.ensure_tables_async(api_keys)
 
-        # Get the full function name as concat of module and qualname, to not accidentally clash
-        full_function_name = func.__module__ + "." + func.__qualname__
-        # But also get the standard function name to easily search for it in database
-        function_name = func.__name__
+            ctx = _build_context(func, args, kwargs, ignore_args, ignore_arg_types)
 
-        # Determine if the function returns or contains Pydantic BaseModel(s)
-        return_type = func.__annotations__.get("return", None)
-        is_pydantic_model = return_type is not None and contains_pydantic_model(
-            return_type
-        )
+            # Fetch cached result in thread pool
+            lookup = await asyncio.to_thread(_fetch_cached, api_keys, ctx, max_age)
 
-        with DBManager(
-            api_keys.sqlalchemy_db_url.get_secret_value()
-        ).get_session() as session:
-            # Try to get cached result
-            statement = (
-                select(FunctionCache)
-                .where(
-                    FunctionCache.function_name == function_name,
-                    FunctionCache.full_function_name == full_function_name,
-                    FunctionCache.args_hash == args_hash,
+            if lookup.hit:
+                logger.debug(
+                    f"{DB_CACHE_LOG_PREFIX} [cache-hit] Cache hit for {ctx.full_function_name}"
                 )
-                .order_by(desc(FunctionCache.created_at))
-            )
-            if max_age is not None:
-                cutoff_time = utcnow() - max_age
-                statement = statement.where(FunctionCache.created_at >= cutoff_time)
-            cached_result = session.exec(statement).first()
-
-        if cached_result:
-            logger.info(
-                # Keep the special [case-hit] identifier so we can easily track it in GCP.
-                f"{DB_CACHE_LOG_PREFIX} [cache-hit] Cache hit for {full_function_name} with args {args_dict} and output {cached_result.result}"
+                return lookup.value
+
+            computed_result = await func(*args, **kwargs)
+            logger.debug(
+                f"{DB_CACHE_LOG_PREFIX} [cache-miss] Cache miss for {ctx.full_function_name}"
             )
-            if is_pydantic_model:
-                # If the output contains any Pydantic models, we need to initialise them.
-                try:
-                    return convert_cached_output_to_pydantic(
-                        return_type, cached_result.result
-                    )
-                except ValueError as e:
-                    # In case of backward-incompatible pydantic model, just treat it as cache miss, to not error out.
-                    logger.warning(
-                        f"{DB_CACHE_LOG_PREFIX} [cache-miss] Can not validate {cached_result=} into {return_type=} because {e=}, treating as cache miss."
+
+            if cache_none or computed_result is not None:
+                # Save cached result in thread pool (fire-and-forget)
+                asyncio.create_task(
+                    asyncio.to_thread(
+                        _save_cached,
+                        api_keys,
+                        ctx,
+                        computed_result,
+                        log_error_on_unsavable_data,
                     )
-                    cached_result = None
-            else:
-                return cached_result.result
+                )
+
+            return computed_result
+
+        return cast(FunctionT, async_wrapper)
+
+    @wraps(func)
+    def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
+        if not api_keys.ENABLE_CACHE:
+            return func(*args, **kwargs)
+
+        # Ensure tables are created before accessing cache
+        _table_manager.ensure_tables_sync(api_keys)
+
+        ctx = _build_context(func, args, kwargs, ignore_args, ignore_arg_types)
+        lookup = _fetch_cached(api_keys, ctx, max_age)
+
+        if lookup.hit:
+            logger.debug(
+                f"{DB_CACHE_LOG_PREFIX} [cache-hit] Cache hit for {ctx.full_function_name}"
+            )
+            return lookup.value
 
-        # On cache miss, compute the result
         computed_result = func(*args, **kwargs)
-        # Keep the special [case-miss] identifier so we can easily track it in GCP.
-        logger.info(
-            f"{DB_CACHE_LOG_PREFIX} [cache-miss] Cache miss for {full_function_name} with args {args_dict}, computed the output {computed_result}"
+        logger.debug(
+            f"{DB_CACHE_LOG_PREFIX} [cache-miss] Cache miss for {ctx.full_function_name}"
         )
 
-        # If postgres access was specified, save it.
         if cache_none or computed_result is not None:
-            cache_entry = FunctionCache(
-                function_name=function_name,
-                full_function_name=full_function_name,
-                args_hash=args_hash,
-                args=args_dict,
-                result=computed_result,
-                created_at=utcnow(),
-            )
-            # Do not raise an exception if saving to the database fails, just log it and let the agent continue the work.
-            try:
-                with DBManager(
-                    api_keys.sqlalchemy_db_url.get_secret_value()
-                ).get_session() as session:
-                    logger.info(
-                        f"{DB_CACHE_LOG_PREFIX} [cache-info] Saving {cache_entry} into database."
-                    )
-                    session.add(cache_entry)
-                    session.commit()
-            except (DataError, psycopg2.errors.UntranslatableCharacter) as e:
-                (logger.error if log_error_on_unsavable_data else logger.warning)(
-                    f"{DB_CACHE_LOG_PREFIX} [cache-error] Failed to save {cache_entry} into database, ignoring, because: {e}"
-                )
-            except Exception:
-                logger.exception(
-                    f"{DB_CACHE_LOG_PREFIX} [cache-error] Failed to save {cache_entry} into database, ignoring."
-                )
+            _save_cached(api_keys, ctx, computed_result, log_error_on_unsavable_data)
 
         return computed_result
 
-    return cast(FunctionT, wrapper)
+    return cast(FunctionT, sync_wrapper)
+
+
+@dataclass
+class CallContext:
+    args_dict: dict[str, Any]
+    args_hash: str
+    function_name: str
+    full_function_name: str
+    return_type: Any
+
+    @property
+    def is_pydantic_model(self) -> bool:
+        return self.return_type is not None and contains_pydantic_model(
+            self.return_type
+        )
+
+
+@dataclass
+class CacheLookup:
+    hit: bool
+    value: Any | None = None
+
+
+def _build_context(
+    func: Callable[..., Any],
+    args: tuple[Any, ...],
+    kwargs: dict[str, Any],
+    ignore_args: Sequence[str] | None,
+    ignore_arg_types: Sequence[type] | None,
+) -> CallContext:
+    signature = inspect.signature(func)
+    bound_arguments = signature.bind(*args, **kwargs)
+    bound_arguments.apply_defaults()
+
+    args_dict: dict[str, Any] = bound_arguments.arguments
+
+    if "self" in args_dict:
+        del args_dict["self"]
+    if "cls" in args_dict:
+        del args_dict["cls"]
+
+    if ignore_args:
+        for arg in ignore_args:
+            if arg in args_dict:
+                del args_dict[arg]
+
+    if ignore_arg_types:
+        args_dict = {
+            k: v
+            for k, v in args_dict.items()
+            if not isinstance(v, tuple(ignore_arg_types))
+        }
+
+    arg_string = json.dumps(args_dict, sort_keys=True, default=str)
+    args_hash = hashlib.md5(arg_string.encode()).hexdigest()
+
+    full_function_name = func.__module__ + "." + func.__qualname__
+    function_name = func.__name__
+
+    # Use get_type_hints to resolve forward references instead of __annotations__
+    try:
+        type_hints = get_type_hints(func)
+        return_type = type_hints.get("return", None)
+    except (NameError, AttributeError, TypeError) as e:
+        # Fallback to raw annotations if get_type_hints fails
+        logger.debug(
+            f"{DB_CACHE_LOG_PREFIX} Failed to resolve type hints for {full_function_name}, falling back to raw annotations: {e}"
+        )
+        return_type = func.__annotations__.get("return", None)
+
+    return CallContext(
+        args_dict=args_dict,
+        args_hash=args_hash,
+        function_name=function_name,
+        full_function_name=full_function_name,
+        return_type=return_type,
+    )
+
+
+def _fetch_cached(
+    api_keys: APIKeys,
+    ctx: CallContext,
+    max_age: timedelta | None,
+) -> CacheLookup:
+    with DBManager(
+        api_keys.sqlalchemy_db_url.get_secret_value()
+    ).get_session() as session:
+        statement = (
+            select(FunctionCache)
+            .where(
+                FunctionCache.function_name == ctx.function_name,
+                FunctionCache.full_function_name == ctx.full_function_name,
+                FunctionCache.args_hash == ctx.args_hash,
+            )
+            .order_by(desc(FunctionCache.created_at))
+        )
+        if max_age is not None:
+            cutoff_time = utcnow() - max_age
+            statement = statement.where(FunctionCache.created_at >= cutoff_time)
+        cached_result = session.exec(statement).first()
+
+    if not cached_result:
+        return CacheLookup(hit=False)
+
+    if ctx.is_pydantic_model:
+        try:
+            value = convert_cached_output_to_pydantic(
+                ctx.return_type, cached_result.result
+            )
+            return CacheLookup(hit=True, value=value)
+        except (ValueError, TypeError) as e:
+            logger.warning(
+                f"{DB_CACHE_LOG_PREFIX} [cache-miss] Failed to validate cached result for {ctx.full_function_name}, treating as cache miss: {e}"
+            )
+            return CacheLookup(hit=False)
+
+    return CacheLookup(hit=True, value=cached_result.result)
+
+
+def _save_cached(
+    api_keys: APIKeys,
+    ctx: CallContext,
+    computed_result: Any,
+    log_error_on_unsavable_data: bool,
+) -> None:
+    cache_entry = FunctionCache(
+        function_name=ctx.function_name,
+        full_function_name=ctx.full_function_name,
+        args_hash=ctx.args_hash,
+        args=ctx.args_dict,
+        result=computed_result,
+        created_at=utcnow(),
+    )
+    try:
+        with DBManager(
+            api_keys.sqlalchemy_db_url.get_secret_value()
+        ).get_session() as session:
+            logger.debug(
+                f"{DB_CACHE_LOG_PREFIX} [cache-save] Saving cache entry for {ctx.full_function_name}"
+            )
+            session.add(cache_entry)
+            session.commit()
+    except (DataError, psycopg2.errors.UntranslatableCharacter) as e:
+        (logger.error if log_error_on_unsavable_data else logger.warning)(
+            f"{DB_CACHE_LOG_PREFIX} [cache-error] Failed to save cache entry for {ctx.full_function_name}: {e}"
+        )
+    except Exception:
+        logger.exception(
+            f"{DB_CACHE_LOG_PREFIX} [cache-error] Failed to save cache entry for {ctx.full_function_name}"
+        )
 
 
 def contains_pydantic_model(return_type: Any) -> bool:
@@ -261,8 +363,8 @@ def convert_cached_output_to_pydantic(return_type: Any, data: Any) -> Any:
         if origin is None:
             if inspect.isclass(return_type) and issubclass(return_type, BaseModel):
                 # Convert the dictionary to a Pydantic model
-                return return_type(
-                    **{
+                return return_type.model_validate(
+                    {
                         k: convert_cached_output_to_pydantic(
                             getattr(return_type, k, None), v
                         )

prediction_market_agent_tooling/tools/caches/serializers.py

@@ -4,6 +4,7 @@ from datetime import date, timedelta
 
 from pydantic import BaseModel
 
+from prediction_market_agent_tooling.gtypes import HexBytes
 from prediction_market_agent_tooling.tools.datetime_utc import DatetimeUTC
 
 
@@ -12,7 +13,7 @@ def json_serializer(x: t.Any) -> str:
 
 
 def json_serializer_default_fn(
-    y: DatetimeUTC | timedelta | date | BaseModel,
+    y: DatetimeUTC | timedelta | date | HexBytes | BaseModel,
 ) -> str | dict[str, t.Any]:
     """
     Used to serialize objects that don't support it by default into a specific string that can be deserialized out later.
@@ -25,8 +26,13 @@ def json_serializer_default_fn(
         return f"timedelta::{y.total_seconds()}"
     elif isinstance(y, date):
         return f"date::{y.isoformat()}"
+    elif isinstance(y, HexBytes):
+        return f"HexBytes::{y.to_0x_hex()}"
     elif isinstance(y, BaseModel):
-        return y.model_dump()
+        # For some reason, Pydantic by default serializes using the field names (not alias),
+        # but also by default, deserializes only using the aliased names.
+        # `by_alias=True` here to work by default with models that have some fields with aliased names.
+        return y.model_dump(by_alias=True)
     raise TypeError(
         f"Unsupported type for the default json serialize function, value is {y}."
     )
@@ -51,6 +57,9 @@ def replace_custom_stringified_objects(obj: t.Any) -> t.Any:
         elif obj.startswith("date::"):
             iso_str = obj[len("date::") :]
             return date.fromisoformat(iso_str)
+        elif obj.startswith("HexBytes::"):
+            hex_str = obj[len("HexBytes::") :]
+            return HexBytes(hex_str)
         else:
             return obj
     elif isinstance(obj, dict):