prediction-market-agent-tooling 0.52.2__py3-none-any.whl → 0.53.0__py3-none-any.whl

prediction_market_agent_tooling/tools/tavily/tavily_models.py

@@ -0,0 +1,84 @@
+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/tools/{tavily_storage/tavily_storage.py → tavily/tavily_search.py}

@@ -4,16 +4,20 @@ import tenacity
 from tavily import TavilyClient
 
 from prediction_market_agent_tooling.config import APIKeys
-from prediction_market_agent_tooling.tools.tavily_storage.tavily_models import (
+from prediction_market_agent_tooling.tools.tavily.tavily_models import (
     TavilyResponse,
-    TavilyStorage,
+    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
 
 
 def tavily_search(
     query: str,
     search_depth: t.Literal["basic", "advanced"] = "advanced",
     topic: t.Literal["general", "news"] = "general",
+    days: int | None = None,
     max_results: int = 5,
     include_domains: t.Sequence[str] | None = None,
     exclude_domains: t.Sequence[str] | None = None,
@@ -35,6 +39,7 @@ def tavily_search(
             search_depth=search_depth,
             topic=topic,
             max_results=max_results,
+            days=days,
             include_domains=include_domains,
             exclude_domains=exclude_domains,
             include_answer=include_answer,
@@ -49,6 +54,7 @@ def tavily_search(
         search_depth=search_depth,
         topic=topic,
         max_results=max_results,
+        days=days,
         include_domains=include_domains,
         exclude_domains=exclude_domains,
         include_answer=include_answer,
@@ -63,6 +69,7 @@ def tavily_search(
             query=query,
             search_depth=search_depth,
             topic=topic,
+            days=days,
             max_results=max_results,
             include_domains=include_domains,
             exclude_domains=exclude_domains,
@@ -80,6 +87,7 @@ def _tavily_search(
     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,
@@ -99,6 +107,7 @@ def _tavily_search(
         query=query,
         search_depth=search_depth,
         topic=topic,
+        days=days,
         max_results=max_results,
         include_domains=include_domains,
         exclude_domains=exclude_domains,
@@ -108,3 +117,20 @@ def _tavily_search(
         use_cache=use_cache,
     )
     return response
+
+
+def get_related_news_since(
+    question: str,
+    days_ago: int,
+    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,
+        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/tools/{tavily_storage/tavily_models.py → tavily/tavily_storage.py}

@@ -2,96 +2,15 @@ import typing as t
 from datetime import timedelta
 
 import tenacity
-from pydantic import BaseModel
-from sqlalchemy import Column
-from sqlalchemy.dialects.postgresql import JSONB
-from sqlmodel import (
-    ARRAY,
-    Field,
-    Session,
-    SQLModel,
-    String,
-    create_engine,
-    desc,
-    select,
-)
+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.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: 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
-    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"],
-        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,
-            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(),
-        )
+from prediction_market_agent_tooling.tools.tavily.tavily_models import (
+    TavilyResponse,
+    TavilyResponseModel,
+)
+from prediction_market_agent_tooling.tools.utils import utcnow
 
 
 class TavilyStorage:
@@ -119,6 +38,7 @@ class TavilyStorage:
         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,
@@ -134,6 +54,7 @@ class TavilyStorage:
             search_depth=search_depth,
             topic=topic,
             max_results=max_results,
+            days=days,
             include_domains=include_domains,
             exclude_domains=exclude_domains,
             include_answer=include_answer,
@@ -152,6 +73,7 @@ class TavilyStorage:
         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,
@@ -167,6 +89,7 @@ class TavilyStorage:
                 .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)

{prediction_market_agent_tooling-0.52.2.dist-info → prediction_market_agent_tooling-0.53.0.dist-info}/METADATA

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

{prediction_market_agent_tooling-0.52.2.dist-info → prediction_market_agent_tooling-0.53.0.dist-info}/RECORD

@@ -88,12 +88,13 @@ prediction_market_agent_tooling/tools/parallelism.py,sha256=6Gou0hbjtMZrYvxjTDFU
 prediction_market_agent_tooling/tools/safe.py,sha256=h0xOO0eNtitClf0fPkn-0oTc6A_bflDTee98V_aiV-A,5195
 prediction_market_agent_tooling/tools/singleton.py,sha256=CiIELUiI-OeS7U7eeHEt0rnVhtQGzwoUdAgn_7u_GBM,729
 prediction_market_agent_tooling/tools/streamlit_user_login.py,sha256=NXEqfjT9Lc9QtliwSGRASIz1opjQ7Btme43H4qJbzgE,3010
-prediction_market_agent_tooling/tools/tavily_storage/tavily_models.py,sha256=99S7w8BvnJRMOnUArGN0g4GVRoG8M0C-XyIFU8HnLn0,6374
-prediction_market_agent_tooling/tools/tavily_storage/tavily_storage.py,sha256=xrtQH9v5pXycBRyc5j45pWqkSffkoc9efNIU1_G633Q,3706
+prediction_market_agent_tooling/tools/tavily/tavily_models.py,sha256=Rz4tZzwCRzPaq49SFT33SCRQrqHXtqWdD9ajb2tGCWc,2723
+prediction_market_agent_tooling/tools/tavily/tavily_search.py,sha256=MK_ozeQbJ014HGiKFPDScjFYq0OGcjY1KPgc9A6qO0M,4511
+prediction_market_agent_tooling/tools/tavily/tavily_storage.py,sha256=t-tZzbCzBBdFedRZDuVBn3A3mIDX8Z5wza6SxWswu_E,4093
 prediction_market_agent_tooling/tools/utils.py,sha256=W-9SqeCKd51BYMRhDjYPQ7lfNO_zE9EvYpmu2r5WXGA,7163
 prediction_market_agent_tooling/tools/web3_utils.py,sha256=dkcjG-LtuaWRh7WEMzRGmZ5B5rsxZTlliFOI6fj-EJ8,11842
-prediction_market_agent_tooling-0.52.2.dist-info/LICENSE,sha256=6or154nLLU6bELzjh0mCreFjt0m2v72zLi3yHE0QbeE,7650
-prediction_market_agent_tooling-0.52.2.dist-info/METADATA,sha256=OKILmVSPUD6h9Z8fyfp1ymB7FJF7WcLMQa5jkj9F-fk,8056
-prediction_market_agent_tooling-0.52.2.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
-prediction_market_agent_tooling-0.52.2.dist-info/entry_points.txt,sha256=m8PukHbeH5g0IAAmOf_1Ahm-sGAMdhSSRQmwtpmi2s8,81
-prediction_market_agent_tooling-0.52.2.dist-info/RECORD,,
+prediction_market_agent_tooling-0.53.0.dist-info/LICENSE,sha256=6or154nLLU6bELzjh0mCreFjt0m2v72zLi3yHE0QbeE,7650
+prediction_market_agent_tooling-0.53.0.dist-info/METADATA,sha256=RrGcyPVn2Gs6Xx9ezoDYzqUMAQfDZVlT63_Uq22GjSc,8056
+prediction_market_agent_tooling-0.53.0.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
+prediction_market_agent_tooling-0.53.0.dist-info/entry_points.txt,sha256=m8PukHbeH5g0IAAmOf_1Ahm-sGAMdhSSRQmwtpmi2s8,81
+prediction_market_agent_tooling-0.53.0.dist-info/RECORD,,