prediction-market-agent-tooling 0.55.2__tar.gz → 0.56.0.dev112__tar.gz

{prediction_market_agent_tooling-0.55.2 → prediction_market_agent_tooling-0.56.0.dev112}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: prediction-market-agent-tooling
-Version: 0.55.2
+Version: 0.56.0.dev112
 Summary: Tools to benchmark, deploy and monitor prediction market agents.
 Author: Gnosis
 Requires-Python: >=3.10,<3.12
@@ -36,6 +36,7 @@ Requires-Dist: psycopg2-binary (>=2.9.9,<3.0.0)
 Requires-Dist: pydantic (>=2.6.1,<3.0.0)
 Requires-Dist: pydantic-settings (>=2.4.0,<3.0.0)
 Requires-Dist: pymongo (>=4.8.0,<5.0.0)
+Requires-Dist: pytest-postgresql (>=6.1.1,<7.0.0)
 Requires-Dist: python-dateutil (>=2.9.0.post0,<3.0.0)
 Requires-Dist: safe-cli (>=1.0.0,<2.0.0)
 Requires-Dist: safe-eth-py (>=6.0.0b41,<7.0.0)

prediction_market_agent_tooling-0.56.0.dev112/prediction_market_agent_tooling/tools/caches/db_cache.py

@@ -0,0 +1,337 @@
+import hashlib
+import inspect
+import json
+from datetime import date, timedelta
+from functools import wraps
+from typing import (
+    Any,
+    Callable,
+    Sequence,
+    TypeVar,
+    cast,
+    get_args,
+    get_origin,
+    overload,
+)
+
+from pydantic import BaseModel
+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.utils import utcnow
+
+FunctionT = TypeVar("FunctionT", bound=Callable[..., Any])
+
+
+class FunctionCache(SQLModel, table=True):
+    __tablename__ = "function_cache"
+    id: int | None = Field(default=None, primary_key=True)
+    function_name: str = Field(index=True)
+    # Args are stored to see what was the function called with.
+    args: Any = Field(sa_column=Column(JSONB, nullable=False))
+    # Args hash is stored as a fast look-up option when looking for cache hits.
+    args_hash: str = Field(index=True)
+    result: Any = Field(sa_column=Column(JSONB, nullable=False))
+    created_at: DatetimeUTC = Field(default_factory=utcnow, index=True)
+
+
+@overload
+def db_cache(
+    func: None = None,
+    *,
+    max_age: timedelta | None = None,
+    cache_none: bool = True,
+    api_keys: APIKeys | None = None,
+    ignore_args: Sequence[str] | None = None,
+    ignore_arg_types: Sequence[type] | None = (APIKeys,),
+) -> Callable[[FunctionT], FunctionT]:
+    ...
+
+
+@overload
+def db_cache(
+    func: FunctionT,
+    *,
+    max_age: timedelta | None = None,
+    cache_none: bool = True,
+    api_keys: APIKeys | None = None,
+    ignore_args: Sequence[str] | None = None,
+    ignore_arg_types: Sequence[type] | None = (APIKeys,),
+) -> FunctionT:
+    ...
+
+
+def db_cache(
+    func: FunctionT | None = None,
+    *,
+    max_age: timedelta | None = None,
+    cache_none: bool = True,
+    api_keys: APIKeys | None = None,
+    ignore_args: Sequence[str] | None = None,
+    ignore_arg_types: Sequence[type]
+    | None = (
+        APIKeys,
+    ),  # Ignore `APIKeys` by default from storing it in DB and using to compute look-up hash, because api keys shouldn't change the computed results.
+) -> FunctionT | Callable[[FunctionT], FunctionT]:
+    if func is None:
+        # Ugly Pythonic way to support this decorator as `@postgres_cache` but also `@postgres_cache(max_age=timedelta(days=3))`
+        def decorator(func: FunctionT) -> FunctionT:
+            return db_cache(
+                func,
+                max_age=max_age,
+                cache_none=cache_none,
+                api_keys=api_keys,
+                ignore_args=ignore_args,
+                ignore_arg_types=ignore_arg_types,
+            )
+
+        return decorator
+
+    api_keys = api_keys if api_keys is not None else APIKeys()
+
+    sqlalchemy_db_url = api_keys.SQLALCHEMY_DB_URL
+    if sqlalchemy_db_url is None:
+        logger.warning(
+            f"SQLALCHEMY_DB_URL not provided in the environment, skipping function caching."
+        )
+
+    engine = (
+        create_engine(
+            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 sqlalchemy_db_url is not None
+        else None
+    )
+
+    # Create table if it doesn't exist
+    if engine is not None:
+        SQLModel.metadata.create_all(engine)
+
+    @wraps(func)
+    def wrapper(*args: Any, **kwargs: Any) -> Any:
+        if not api_keys.ENABLE_CACHE:
+            return 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 = convert_pydantic_to_dict(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))
+            }
+
+        # 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()
+
+        # Get the function name as concat of module and qualname, to not accidentally clash
+        function_name = func.__module__ + "." + func.__qualname__
+
+        # Determine if the function returns or contains Pydantic BaseModel(s)
+        return_type = func.__annotations__.get("return", None)
+        is_pydantic_model = False
+
+        if return_type is not None and contains_pydantic_model(return_type):
+            is_pydantic_model = True
+
+        # If postgres access was specified, try to find a hit
+        if engine is not None:
+            with Session(engine) as session:
+                # Try to get cached result
+                statement = (
+                    select(FunctionCache)
+                    .where(
+                        FunctionCache.function_name == function_name,
+                        FunctionCache.args_hash == 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()
+        else:
+            cached_result = None
+
+        if cached_result:
+            logger.info(
+                f"Cache hit for {function_name} with args {args_dict} and output {cached_result.result}"
+            )
+            if is_pydantic_model:
+                try:
+                    return convert_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"Can not validate {cached_result=} into {return_type=} because {e=}, treating as cache miss."
+                    )
+                    cached_result = None
+            else:
+                return cached_result.result
+
+        # On cache miss, compute the result
+        computed_result = func(*args, **kwargs)
+        logger.info(
+            f"Cache miss for {function_name} with args {args_dict}, computed the output {computed_result}"
+        )
+
+        # If postgres access was specified, save it to dB.
+        if engine is not None and (cache_none or computed_result is not None):
+            # Call the original function
+            result_data = (
+                convert_pydantic_to_dict(computed_result)
+                if is_pydantic_model
+                else computed_result
+            )
+            cache_entry = FunctionCache(
+                function_name=function_name,
+                args_hash=args_hash,
+                args=args_dict,
+                result=result_data,
+                created_at=utcnow(),
+            )
+            with Session(engine) as session:
+                logger.info(f"Saving {cache_entry} into database.")
+                session.add(cache_entry)
+                session.commit()
+
+        return computed_result
+
+    return cast(FunctionT, wrapper)
+
+
+def contains_pydantic_model(tp: Any) -> bool:
+    if tp is None:
+        return False
+    origin = get_origin(tp)
+    if origin is not None:
+        return any(contains_pydantic_model(arg) for arg in get_args(tp))
+    if inspect.isclass(tp):
+        return issubclass(tp, BaseModel)
+    return False
+
+
+def json_serializer_default_fn(y: Any) -> Any:
+    if isinstance(y, DatetimeUTC):
+        return f"DatetimeUTC::{y.isoformat()}"
+    elif isinstance(y, timedelta):
+        return f"timedelta::{y.total_seconds()}"
+    elif isinstance(y, date):
+        return f"date::{y.isoformat()}"
+    raise TypeError(
+        f"Unsuported type for the default json serialize function, value is {y}."
+    )
+
+
+def json_serializer(x: Any) -> str:
+    return json.dumps(x, default=json_serializer_default_fn)
+
+
+def replace_custom_stringified_objects(obj: Any) -> Any:
+    if isinstance(obj, str):
+        if obj.startswith("DatetimeUTC::"):
+            iso_str = obj[len("DatetimeUTC::") :]
+            return DatetimeUTC.to_datetime_utc(iso_str)
+        elif obj.startswith("timedelta::"):
+            total_seconds_str = obj[len("timedelta::") :]
+            return timedelta(seconds=float(total_seconds_str))
+        elif obj.startswith("date::"):
+            iso_str = obj[len("date::") :]
+            return date.fromisoformat(iso_str)
+        else:
+            return obj
+    elif isinstance(obj, dict):
+        return {k: replace_custom_stringified_objects(v) for k, v in obj.items()}
+    elif isinstance(obj, list):
+        return [replace_custom_stringified_objects(item) for item in obj]
+    else:
+        return obj
+
+
+def json_deserializer(s: str) -> Any:
+    data = json.loads(s)
+    return replace_custom_stringified_objects(data)
+
+
+def convert_pydantic_to_dict(value: Any) -> Any:
+    if isinstance(value, BaseModel):
+        return value.model_dump()
+    elif isinstance(value, dict):
+        return {k: convert_pydantic_to_dict(v) for k, v in value.items()}
+    elif isinstance(value, (list, tuple)):
+        return type(value)(convert_pydantic_to_dict(v) for v in value)
+    elif isinstance(value, set):
+        return {convert_pydantic_to_dict(v) for v in value}
+    else:
+        return value
+
+
+def convert_to_pydantic(model: Any, data: Any) -> Any:
+    # Get the origin and arguments of the model type
+    origin = get_origin(model)
+    args = get_args(model)
+
+    # Check if the data is a dictionary
+    if isinstance(data, dict):
+        # If the model has no origin, check if it is a subclass of BaseModel
+        if origin is None:
+            if inspect.isclass(model) and issubclass(model, BaseModel):
+                # Convert the dictionary to a Pydantic model
+                return model(
+                    **{
+                        k: convert_to_pydantic(getattr(model, k, None), v)
+                        for k, v in data.items()
+                    }
+                )
+            else:
+                # If not a Pydantic model, return the data as is
+                return data
+        # If the origin is a dictionary, convert keys and values
+        elif origin is dict:
+            key_type, value_type = args
+            return {
+                convert_to_pydantic(key_type, k): convert_to_pydantic(value_type, v)
+                for k, v in data.items()
+            }
+        else:
+            # If the origin is not a dictionary, return the data as is
+            return data
+    # Check if the data is a list
+    elif isinstance(data, list):
+        # If the origin is a list, convert each item
+        if origin is list:
+            item_type = args[0]
+            return [convert_to_pydantic(item_type, item) for item in data]
+        else:
+            # If the origin is not a list, return the data as is
+            return data
+    else:
+        # If the data is neither a dictionary nor a list, return it as is
+        return data

{prediction_market_agent_tooling-0.55.2 → prediction_market_agent_tooling-0.56.0.dev112}/prediction_market_agent_tooling/tools/google.py

@@ -5,7 +5,7 @@ from googleapiclient.discovery import build
 
 from prediction_market_agent_tooling.config import APIKeys
 from prediction_market_agent_tooling.loggers import logger
-from prediction_market_agent_tooling.tools.cache import persistent_inmemory_cache
+from prediction_market_agent_tooling.tools.caches.db_cache import db_cache
 
 
 @tenacity.retry(
@@ -13,7 +13,7 @@ from prediction_market_agent_tooling.tools.cache import persistent_inmemory_cach
     stop=tenacity.stop_after_attempt(3),
     after=lambda x: logger.debug(f"search_google failed, {x.attempt_number=}."),
 )
-@persistent_inmemory_cache
+@db_cache
 def search_google(
     query: str | None = None,
     num: int = 3,

{prediction_market_agent_tooling-0.55.2 → prediction_market_agent_tooling-0.56.0.dev112}/prediction_market_agent_tooling/tools/is_invalid.py

@@ -2,7 +2,7 @@ import tenacity
 
 from prediction_market_agent_tooling.config import APIKeys
 from prediction_market_agent_tooling.loggers import logger
-from prediction_market_agent_tooling.tools.cache import persistent_inmemory_cache
+from prediction_market_agent_tooling.tools.caches.db_cache import db_cache
 from prediction_market_agent_tooling.tools.is_predictable import (
     parse_decision_yes_no_completion,
 )
@@ -54,9 +54,9 @@ Finally, write your final decision, write `decision: ` followed by either "yes i
 """
 
 
-@persistent_inmemory_cache
 @tenacity.retry(stop=tenacity.stop_after_attempt(3), wait=tenacity.wait_fixed(1))
 @observe()
+@db_cache
 def is_invalid(
     question: str,
     engine: str = "gpt-4o",

{prediction_market_agent_tooling-0.55.2 → prediction_market_agent_tooling-0.56.0.dev112}/prediction_market_agent_tooling/tools/is_predictable.py

@@ -2,7 +2,7 @@ import tenacity
 
 from prediction_market_agent_tooling.config import APIKeys
 from prediction_market_agent_tooling.loggers import logger
-from prediction_market_agent_tooling.tools.cache import persistent_inmemory_cache
+from prediction_market_agent_tooling.tools.caches.db_cache import db_cache
 from prediction_market_agent_tooling.tools.langfuse_ import (
     get_langfuse_langchain_config,
     observe,
@@ -76,9 +76,9 @@ Finally, write your final decision, write `decision: ` followed by either "yes i
 """
 
 
-@persistent_inmemory_cache
 @tenacity.retry(stop=tenacity.stop_after_attempt(3), wait=tenacity.wait_fixed(1))
 @observe()
+@db_cache
 def is_predictable_binary(
     question: str,
     engine: str = "gpt-4-1106-preview",
@@ -112,9 +112,9 @@ def is_predictable_binary(
     return parse_decision_yes_no_completion(question, completion)
 
 
-@persistent_inmemory_cache
 @tenacity.retry(stop=tenacity.stop_after_attempt(3), wait=tenacity.wait_fixed(1))
 @observe()
+@db_cache
 def is_predictable_without_description(
     question: str,
     description: str,

{prediction_market_agent_tooling-0.55.2 → prediction_market_agent_tooling-0.56.0.dev112}/prediction_market_agent_tooling/tools/relevant_news_analysis/relevant_news_analysis.py

@@ -1,4 +1,4 @@
-from datetime import datetime, timedelta
+from datetime import date, timedelta
 
 from langchain_core.output_parsers import PydanticOutputParser
 from langchain_core.prompts import PromptTemplate
@@ -20,8 +20,7 @@ from prediction_market_agent_tooling.tools.relevant_news_analysis.relevant_news_
 from prediction_market_agent_tooling.tools.tavily.tavily_search import (
     get_relevant_news_since,
 )
-from prediction_market_agent_tooling.tools.tavily.tavily_storage import TavilyStorage
-from prediction_market_agent_tooling.tools.utils import check_not_none, utcnow
+from prediction_market_agent_tooling.tools.utils import check_not_none
 
 SUMMARISE_RELEVANT_NEWS_PROMPT_TEMPLATE = """
 You are an expert news analyst, tracking stories that may affect your prediction to the outcome of a particular QUESTION.
@@ -55,7 +54,7 @@ For your analysis, you should:
 def analyse_news_relevance(
     raw_content: str,
     question: str,
-    date_of_interest: datetime,
+    date_of_interest: date,
     model: str,
     temperature: float,
 ) -> RelevantNewsAnalysis:
@@ -91,19 +90,18 @@ def analyse_news_relevance(
 def get_certified_relevant_news_since(
     question: str,
     days_ago: int,
-    tavily_storage: TavilyStorage | None = None,
 ) -> RelevantNews | None:
     """
     Get relevant news since a given date for a given question. Retrieves
     possibly relevant news from tavily, then checks that it is relevant via
     an LLM call.
     """
+    news_since = date.today() - timedelta(days=days_ago)
     results = get_relevant_news_since(
         question=question,
-        days_ago=days_ago,
+        news_since=news_since,
         score_threshold=0.0,  # Be conservative to avoid missing relevant information
         max_results=3,  # A tradeoff between cost and quality. 3 seems to be a good balance.
-        tavily_storage=tavily_storage,
     )
 
     # Sort results by descending 'relevance score' to maximise the chance of
@@ -118,7 +116,7 @@ def get_certified_relevant_news_since(
         relevant_news_analysis = analyse_news_relevance(
             raw_content=check_not_none(result.raw_content),
             question=question,
-            date_of_interest=utcnow() - timedelta(days=days_ago),
+            date_of_interest=news_since,
             model="gpt-4o",  # 4o-mini isn't good enough, 1o and 1o-mini are too expensive
             temperature=0.0,
         )
@@ -140,7 +138,6 @@ def get_certified_relevant_news_since_cached(
     question: str,
     days_ago: int,
     cache: RelevantNewsResponseCache,
-    tavily_storage: TavilyStorage | None = None,
 ) -> RelevantNews | None:
     cached = cache.find(question=question, days_ago=days_ago)
 
@@ -150,7 +147,6 @@ def get_certified_relevant_news_since_cached(
         relevant_news = get_certified_relevant_news_since(
             question=question,
             days_ago=days_ago,
-            tavily_storage=tavily_storage,
         )
         cache.save(
             question=question,

prediction_market_agent_tooling-0.56.0.dev112/prediction_market_agent_tooling/tools/tavily/tavily_models.py

@@ -0,0 +1,18 @@
+from pydantic import BaseModel
+
+
+class TavilyResult(BaseModel):
+    title: str
+    url: str
+    content: str
+    score: float
+    raw_content: str | None
+
+
+class TavilyResponse(BaseModel):
+    query: str
+    follow_up_questions: str | None = None
+    answer: str
+    images: list[str]
+    results: list[TavilyResult]
+    response_time: float

{prediction_market_agent_tooling-0.55.2 → prediction_market_agent_tooling-0.56.0.dev112}/prediction_market_agent_tooling/tools/tavily/tavily_search.py

@@ -1,23 +1,25 @@
 import typing as t
+from datetime import date
 
 import tenacity
 from tavily import TavilyClient
 
 from prediction_market_agent_tooling.config import APIKeys
+from prediction_market_agent_tooling.tools.caches.db_cache import db_cache
 from prediction_market_agent_tooling.tools.tavily.tavily_models import (
     TavilyResponse,
     TavilyResult,
 )
-from prediction_market_agent_tooling.tools.tavily.tavily_storage import TavilyStorage
 
 DEFAULT_SCORE_THRESHOLD = 0.75  # Based on some empirical testing, anything lower wasn't very relevant to the question being asked
 
 
+@db_cache
 def tavily_search(
     query: str,
     search_depth: t.Literal["basic", "advanced"] = "advanced",
     topic: t.Literal["general", "news"] = "general",
-    days: int | None = None,
+    news_since: date | None = None,
     max_results: int = 5,
     include_domains: t.Sequence[str] | None = None,
     exclude_domains: t.Sequence[str] | None = None,
@@ -26,34 +28,16 @@ def tavily_search(
     include_images: bool = True,
     use_cache: bool = False,
     api_keys: APIKeys | None = None,
-    tavily_storage: TavilyStorage | None = None,
 ) -> TavilyResponse:
     """
-    Wrapper around Tavily's search method that will save the response to `TavilyResponseCache`, if provided.
-
     Argument default values are different from the original method, to return everything by default, because it can be handy in the future and it doesn't increase the costs.
     """
-    if topic == "news" and days is None:
-        raise ValueError("When topic is 'news', days must be an integer")
-    if topic == "general" and days is not None:
-        raise ValueError("When topic is 'general', days must be None")
+    if topic == "news" and news_since is None:
+        raise ValueError("When topic is 'news', news_since must be provided")
+    if topic == "general" and news_since is not None:
+        raise ValueError("When topic is 'general', news_since must be None")
 
-    if tavily_storage and (
-        response_parsed := tavily_storage.find(
-            query=query,
-            search_depth=search_depth,
-            topic=topic,
-            max_results=max_results,
-            days=days,
-            include_domains=include_domains,
-            exclude_domains=exclude_domains,
-            include_answer=include_answer,
-            include_raw_content=include_raw_content,
-            include_images=include_images,
-            use_cache=use_cache,
-        )
-    ):
-        return response_parsed
+    days = None if news_since is None else (date.today() - news_since).days
     response = _tavily_search(
         query=query,
         search_depth=search_depth,
@@ -69,21 +53,7 @@ def tavily_search(
         api_keys=api_keys,
     )
     response_parsed = TavilyResponse.model_validate(response)
-    if tavily_storage:
-        tavily_storage.save(
-            query=query,
-            search_depth=search_depth,
-            topic=topic,
-            days=days,
-            max_results=max_results,
-            include_domains=include_domains,
-            exclude_domains=exclude_domains,
-            include_answer=include_answer,
-            include_raw_content=include_raw_content,
-            include_images=include_images,
-            use_cache=use_cache,
-            response=response_parsed,
-        )
+
     return response_parsed
 
 
@@ -131,16 +101,14 @@ def _tavily_search(
 
 def get_relevant_news_since(
     question: str,
-    days_ago: int,
+    news_since: date,
     score_threshold: float = DEFAULT_SCORE_THRESHOLD,
     max_results: int = 3,
-    tavily_storage: TavilyStorage | None = None,
 ) -> list[TavilyResult]:
     news = tavily_search(
         query=question,
-        days=days_ago,
+        news_since=news_since,
         max_results=max_results,
         topic="news",
-        tavily_storage=tavily_storage,
     )
     return [r for r in news.results if r.score > score_threshold]

{prediction_market_agent_tooling-0.55.2 → prediction_market_agent_tooling-0.56.0.dev112}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "prediction-market-agent-tooling"
-version = "0.55.2"
+version = "0.56.0.dev112"
 description = "Tools to benchmark, deploy and monitor prediction market agents."
 authors = ["Gnosis"]
 readme = "README.md"
@@ -52,6 +52,7 @@ python-dateutil = "^2.9.0.post0"
 types-python-dateutil = "^2.9.0.20240906"
 pinatapy-vourhey = "^0.2.0"
 hishel = "^0.0.31"
+pytest-postgresql = "^6.1.1"
 
 [tool.poetry.extras]
 openai = ["openai"]

prediction_market_agent_tooling-0.55.2/prediction_market_agent_tooling/tools/tavily/tavily_models.py

@@ -1,84 +0,0 @@
-import typing as t
-
-from pydantic import BaseModel
-from sqlalchemy import Column
-from sqlalchemy.dialects.postgresql import JSONB
-from sqlmodel import ARRAY, Field, SQLModel, String
-
-from prediction_market_agent_tooling.tools.utils import DatetimeUTC, utcnow
-
-
-class TavilyResult(BaseModel):
-    title: str
-    url: str
-    content: str
-    score: float
-    raw_content: str | None
-
-
-class TavilyResponse(BaseModel):
-    query: str
-    follow_up_questions: str | None = None
-    answer: str
-    images: list[str]
-    results: list[TavilyResult]
-    response_time: float
-
-
-class TavilyResponseModel(SQLModel, table=True):
-    __tablename__ = "tavily_response"
-    __table_args__ = {"extend_existing": True}
-    id: int | None = Field(None, primary_key=True)
-    agent_id: str = Field(index=True, nullable=False)
-    # Parameters used to execute the search
-    query: str = Field(index=True, nullable=False)
-    search_depth: str
-    topic: str
-    days: int | None = Field(default=None, nullable=True)
-    max_results: int
-    include_domains: list[str] | None = Field(
-        None, sa_column=Column(ARRAY(String), nullable=True)
-    )
-    exclude_domains: list[str] | None = Field(
-        None, sa_column=Column(ARRAY(String), nullable=True)
-    )
-    include_answer: bool
-    include_raw_content: bool
-    include_images: bool
-    use_cache: bool
-    # Datetime at the time of search response and response from the search
-    datetime_: DatetimeUTC = Field(index=True, nullable=False)
-    response: dict[str, t.Any] = Field(sa_column=Column(JSONB, nullable=False))
-
-    @staticmethod
-    def from_model(
-        agent_id: str,
-        query: str,
-        search_depth: t.Literal["basic", "advanced"],
-        topic: t.Literal["general", "news"],
-        days: int | None,
-        max_results: int,
-        include_domains: t.Sequence[str] | None,
-        exclude_domains: t.Sequence[str] | None,
-        include_answer: bool,
-        include_raw_content: bool,
-        include_images: bool,
-        use_cache: bool,
-        response: TavilyResponse,
-    ) -> "TavilyResponseModel":
-        return TavilyResponseModel(
-            agent_id=agent_id,
-            query=query,
-            search_depth=search_depth,
-            topic=topic,
-            days=days,
-            max_results=max_results,
-            include_domains=sorted(include_domains) if include_domains else None,
-            exclude_domains=sorted(exclude_domains) if exclude_domains else None,
-            include_answer=include_answer,
-            include_raw_content=include_raw_content,
-            include_images=include_images,
-            use_cache=use_cache,
-            datetime_=utcnow(),
-            response=response.model_dump(),
-        )

prediction_market_agent_tooling-0.55.2/prediction_market_agent_tooling/tools/tavily/tavily_storage.py

@@ -1,105 +0,0 @@
-import typing as t
-from datetime import timedelta
-
-import tenacity
-from sqlmodel import 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.tavily.tavily_models import (
-    TavilyResponse,
-    TavilyResponseModel,
-)
-from prediction_market_agent_tooling.tools.utils import utcnow
-
-
-class TavilyStorage:
-    def __init__(self, agent_id: str, sqlalchemy_db_url: str | None = None):
-        self.agent_id = agent_id
-        self.engine = create_engine(
-            sqlalchemy_db_url
-            if sqlalchemy_db_url
-            else APIKeys().sqlalchemy_db_url.get_secret_value()
-        )
-        self._initialize_db()
-
-    def _initialize_db(self) -> None:
-        """
-        Creates the tables if they don't exist
-        """
-
-        # trick for making models import mandatory - models must be imported for metadata.create_all to work
-        logger.debug(f"tables being added {TavilyResponseModel}")
-        SQLModel.metadata.create_all(self.engine)
-
-    @tenacity.retry(stop=tenacity.stop_after_attempt(3), wait=tenacity.wait_fixed(1))
-    def save(
-        self,
-        query: str,
-        search_depth: t.Literal["basic", "advanced"],
-        topic: t.Literal["general", "news"],
-        days: int | None,
-        max_results: int,
-        include_domains: t.Sequence[str] | None,
-        exclude_domains: t.Sequence[str] | None,
-        include_answer: bool,
-        include_raw_content: bool,
-        include_images: bool,
-        use_cache: bool,
-        response: TavilyResponse,
-    ) -> None:
-        db_item = TavilyResponseModel.from_model(
-            agent_id=self.agent_id,
-            query=query,
-            search_depth=search_depth,
-            topic=topic,
-            max_results=max_results,
-            days=days,
-            include_domains=include_domains,
-            exclude_domains=exclude_domains,
-            include_answer=include_answer,
-            include_raw_content=include_raw_content,
-            include_images=include_images,
-            use_cache=use_cache,
-            response=response,
-        )
-        with Session(self.engine) as session:
-            session.add(db_item)
-            session.commit()
-
-    @tenacity.retry(stop=tenacity.stop_after_attempt(3), wait=tenacity.wait_fixed(1))
-    def find(
-        self,
-        query: str,
-        search_depth: t.Literal["basic", "advanced"],
-        topic: t.Literal["general", "news"],
-        days: int | None,
-        max_results: int,
-        include_domains: t.Sequence[str] | None,
-        exclude_domains: t.Sequence[str] | None,
-        include_answer: bool,
-        include_raw_content: bool,
-        include_images: bool,
-        use_cache: bool,
-        max_age: timedelta = timedelta(days=1),
-    ) -> TavilyResponse | None:
-        with Session(self.engine) as session:
-            sql_query = (
-                select(TavilyResponseModel)
-                .where(TavilyResponseModel.query == query)
-                .where(TavilyResponseModel.search_depth == search_depth)
-                .where(TavilyResponseModel.topic == topic)
-                .where(TavilyResponseModel.days == days)
-                .where(TavilyResponseModel.max_results == max_results)
-                .where(TavilyResponseModel.include_domains == include_domains)
-                .where(TavilyResponseModel.exclude_domains == exclude_domains)
-                .where(TavilyResponseModel.include_answer == include_answer)
-                .where(TavilyResponseModel.include_raw_content == include_raw_content)
-                .where(TavilyResponseModel.include_images == include_images)
-                .where(TavilyResponseModel.use_cache == use_cache)
-                .where(TavilyResponseModel.datetime_ >= utcnow() - max_age)
-            )
-            item = session.exec(
-                sql_query.order_by(desc(TavilyResponseModel.datetime_))
-            ).first()
-            return TavilyResponse.model_validate(item.response) if item else None