prediction-market-agent-tooling 0.69.8__py3-none-any.whl → 0.69.9__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,
 )
 
@@ -101,136 +104,235 @@ 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)
+    # Check if the decorated function is async
+    if inspect.iscoroutinefunction(func):
 
-        DBManager(api_keys.sqlalchemy_db_url.get_secret_value()).create_tables(
-            [FunctionCache]
-        )
+        @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)
 
-        # 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))
-            }
+            # Run blocking database operations in thread pool
 
-        # 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 in thread pool
+            await asyncio.to_thread(_ensure_tables, 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(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 _ensure_tables(api_keys: APIKeys) -> None:
+    DBManager(api_keys.sqlalchemy_db_url.get_secret_value()).create_tables(
+        [FunctionCache]
+    )
+
+
+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:

{prediction_market_agent_tooling-0.69.8.dist-info → prediction_market_agent_tooling-0.69.9.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: prediction-market-agent-tooling
-Version: 0.69.8
+Version: 0.69.9
 Summary: Tools to benchmark, deploy and monitor prediction market agents.
 License-File: LICENSE
 Author: Gnosis

{prediction_market_agent_tooling-0.69.8.dist-info → prediction_market_agent_tooling-0.69.9.dist-info}/RECORD

@@ -89,7 +89,7 @@ prediction_market_agent_tooling/tools/balances.py,sha256=Osab21btfJDw2Y-jT_TV-KH
 prediction_market_agent_tooling/tools/betting_strategies/kelly_criterion.py,sha256=o5ba633gKiDqV4t_C2d9FWwH-HkRAOZd8FcZTYvbj6g,14451
 prediction_market_agent_tooling/tools/betting_strategies/stretch_bet_between.py,sha256=THMXwFlskvzbjnX_OiYtDSzI8XVFyULWfP2525_9UGc,429
 prediction_market_agent_tooling/tools/betting_strategies/utils.py,sha256=MpS3FOMn0C7nbmbQRUT9QwSh3UzzsgGczP91iSMr9wo,261
-prediction_market_agent_tooling/tools/caches/db_cache.py,sha256=rZIGhgijquwwPtp_qncSAPR1SDF2XxIVZL1ir0fgzWw,12127
+prediction_market_agent_tooling/tools/caches/db_cache.py,sha256=V6o6UdesjkKzSJMhqkUtD76cJGPaNhuwA4OL2chIYSI,13801
 prediction_market_agent_tooling/tools/caches/inmemory_cache.py,sha256=ZW5iI5rmjqeAebu5T7ftRnlkxiL02IC-MxCfDB80x7w,1506
 prediction_market_agent_tooling/tools/caches/serializers.py,sha256=vFDx4fsPxclXp2q0sv27j4al_M_Tj9aR2JJP-xNHQXA,2151
 prediction_market_agent_tooling/tools/contract.py,sha256=BzpAFcbKl_KqwgAlaXx63Fg8jzr0EO3qEeOs1K11CPA,33905
@@ -137,8 +137,8 @@ prediction_market_agent_tooling/tools/tokens/usd.py,sha256=DPO-4HBTy1-TZHKL_9CnH
 prediction_market_agent_tooling/tools/transaction_cache.py,sha256=K5YKNL2_tR10Iw2TD9fuP-CTGpBbZtNdgbd0B_R7pjg,1814
 prediction_market_agent_tooling/tools/utils.py,sha256=ruq6P5TFs8CBHxeBLj1Plpx7kuNFPpDgMsJGQgDiRNs,8785
 prediction_market_agent_tooling/tools/web3_utils.py,sha256=CDbaidlLeQ4VHzSg150L7QNfHfGveljSePGuDVFEYqc,13963
-prediction_market_agent_tooling-0.69.8.dist-info/METADATA,sha256=bhMWQsWioHf_rG2CIgiOkVBPCqD7B5XFfibk1BFpgUs,8890
-prediction_market_agent_tooling-0.69.8.dist-info/WHEEL,sha256=M5asmiAlL6HEcOq52Yi5mmk9KmTVjY2RDPtO4p9DMrc,88
-prediction_market_agent_tooling-0.69.8.dist-info/entry_points.txt,sha256=m8PukHbeH5g0IAAmOf_1Ahm-sGAMdhSSRQmwtpmi2s8,81
-prediction_market_agent_tooling-0.69.8.dist-info/licenses/LICENSE,sha256=6or154nLLU6bELzjh0mCreFjt0m2v72zLi3yHE0QbeE,7650
-prediction_market_agent_tooling-0.69.8.dist-info/RECORD,,
+prediction_market_agent_tooling-0.69.9.dist-info/METADATA,sha256=fBezbmaBxLVzpMNtCsmCnE2Pc8ywbXJTnCRF7HMJcYw,8890
+prediction_market_agent_tooling-0.69.9.dist-info/WHEEL,sha256=M5asmiAlL6HEcOq52Yi5mmk9KmTVjY2RDPtO4p9DMrc,88
+prediction_market_agent_tooling-0.69.9.dist-info/entry_points.txt,sha256=m8PukHbeH5g0IAAmOf_1Ahm-sGAMdhSSRQmwtpmi2s8,81
+prediction_market_agent_tooling-0.69.9.dist-info/licenses/LICENSE,sha256=6or154nLLU6bELzjh0mCreFjt0m2v72zLi3yHE0QbeE,7650
+prediction_market_agent_tooling-0.69.9.dist-info/RECORD,,