prediction-market-agent-tooling 0.56.2.dev129__tar.gz → 0.56.2.dev137__tar.gz

{prediction_market_agent_tooling-0.56.2.dev129 → prediction_market_agent_tooling-0.56.2.dev137}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: prediction-market-agent-tooling
-Version: 0.56.2.dev129
+Version: 0.56.2.dev137
 Summary: Tools to benchmark, deploy and monitor prediction market agents.
 Author: Gnosis
 Requires-Python: >=3.10,<3.12

{prediction_market_agent_tooling-0.56.2.dev129 → prediction_market_agent_tooling-0.56.2.dev137}/prediction_market_agent_tooling/tools/caches/db_cache.py

@@ -2,7 +2,7 @@ import hashlib
 import inspect
 import json
 from datetime import date, timedelta
-from functools import wraps
+from functools import partial, wraps
 from typing import (
     Any,
     Callable,
@@ -15,13 +15,14 @@ from typing import (
 )
 
 from pydantic import BaseModel
-from sqlalchemy import Column, Engine
+from sqlalchemy import Column
 from sqlalchemy.dialects.postgresql import JSONB
 from sqlmodel import Field, Session, SQLModel, create_engine, 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.pickle_utils import InitialiseNonPickable
 from prediction_market_agent_tooling.tools.utils import utcnow
 
 FunctionT = TypeVar("FunctionT", bound=Callable[..., Any])
@@ -40,9 +41,6 @@ class FunctionCache(SQLModel, table=True):
     created_at: DatetimeUTC = Field(default_factory=utcnow, index=True)
 
 
-DB_CACHE_ENGINE: None | Engine = None
-
-
 @overload
 def db_cache(
     func: None = None,
@@ -93,6 +91,18 @@ def db_cache(
         return decorator
 
     api_keys = api_keys if api_keys is not None else APIKeys()
+    wrapped_engine = InitialiseNonPickable(
+        partial(
+            create_engine,
+            api_keys.sqlalchemy_db_url.get_secret_value(),
+            # Use custom json serializer and deserializer, because otherwise, for example `datetime` serialization would fail.
+            json_serializer=json_serializer,
+            json_deserializer=json_deserializer,
+        )
+    )
+
+    if api_keys.ENABLE_CACHE:
+        SQLModel.metadata.create_all(wrapped_engine.get_value())
 
     @wraps(func)
     def wrapper(*args: Any, **kwargs: Any) -> Any:
@@ -100,16 +110,7 @@ def db_cache(
         if not api_keys.ENABLE_CACHE:
             return func(*args, **kwargs)
 
-        global DB_CACHE_ENGINE
-        if DB_CACHE_ENGINE is None:
-            DB_CACHE_ENGINE = create_engine(
-                api_keys.sqlalchemy_db_url.get_secret_value(),
-                # Use custom json serializer and deserializer, because otherwise, for example `datetime` serialization would fail.
-                json_serializer=json_serializer,
-                json_deserializer=json_deserializer,
-            )
-            # Create table if it doesn't exist
-            SQLModel.metadata.create_all(DB_CACHE_ENGINE)
+        engine = wrapped_engine.get_value()
 
         # Convert *args and **kwargs to a single dictionary, where we have names for arguments passed as args as well.
         signature = inspect.signature(func)
@@ -155,7 +156,7 @@ def db_cache(
         if return_type is not None and contains_pydantic_model(return_type):
             is_pydantic_model = True
 
-        with Session(DB_CACHE_ENGINE) as session:
+        with Session(engine) as session:
             # Try to get cached result
             statement = (
                 select(FunctionCache)
@@ -208,7 +209,7 @@ def db_cache(
                 result=computed_result,
                 created_at=utcnow(),
             )
-            with Session(DB_CACHE_ENGINE) as session:
+            with Session(engine) as session:
                 logger.info(f"Saving {cache_entry} into database.")
                 session.add(cache_entry)
                 session.commit()

prediction_market_agent_tooling-0.56.2.dev137/prediction_market_agent_tooling/tools/pickle_utils.py

@@ -0,0 +1,31 @@
+import typing as t
+
+InitialisedValue = t.TypeVar("InitialisedValue")
+
+
+class InitialiseNonPickable(t.Generic[InitialisedValue]):
+    """
+    Use this class to wrap values that you want to be shared within a thread,
+    but they are re-initialised for a new processes.
+
+    Initialiser for the value still needs to be pickable.
+    """
+
+    def __init__(self, initialiser: t.Callable[[], InitialisedValue]) -> None:
+        self.value: InitialisedValue | None = None
+        self.initialiser = initialiser
+
+    def __getstate__(self) -> dict[str, t.Any]:
+        # During pickling, always return `value` as just None, which is pickable and this class will re-initialise it in `get_value` when called.
+        return {"value": None, "initialiser": self.initialiser}
+
+    def __setstate__(self, d: dict[str, t.Any]) -> None:
+        self.value = d["value"]
+        self.initialiser = d["initialiser"]
+
+    def get_value(self) -> InitialisedValue:
+        """Use this function to get the wrapped value, which will be initialised if necessary."""
+        if self.value is None:
+            self.value = self.initialiser()
+
+        return self.value

{prediction_market_agent_tooling-0.56.2.dev129 → prediction_market_agent_tooling-0.56.2.dev137}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "prediction-market-agent-tooling"
-version = "0.56.2.dev129"
+version = "0.56.2.dev137"
 description = "Tools to benchmark, deploy and monitor prediction market agents."
 authors = ["Gnosis"]
 readme = "README.md"