Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Retry tavily outside of model parsing and allow to not generate subqueries #104

Merged
merged 3 commits into from
Aug 1, 2024
Merged

Retry tavily outside of model parsing and allow to not generate subqueries #104

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
ac609c2
Retry tavily outside of model parsing
kongzii Aug 1, 2024
ecc58b8
Allow to not generate subqueries
kongzii Aug 1, 2024
0639c36
move the condition
kongzii Aug 1, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions prediction_prophet/functions/generate_subqueries.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,9 @@
Limit your searches to {search_limit}.
"""
def generate_subqueries(query: str, limit: int, model: str, api_key: SecretStr | None = None) -> list[str]:
if limit == 0:
return [query]

if api_key == None:
api_key = secret_str_from_env("OPENAI_API_KEY")

Expand Down
34 changes: 21 additions & 13 deletions prediction_prophet/functions/web_search.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
import os
import tenacity
import typing as t
from tavily import TavilyClient
from pydantic.types import SecretStr

Expand All @@ -8,19 +8,9 @@
from prediction_prophet.functions.cache import persistent_inmemory_cache


@tenacity.retry(stop=tenacity.stop_after_attempt(3), wait=tenacity.wait_fixed(1), reraise=True)
@persistent_inmemory_cache

def web_search(query: str, max_results: int = 5, tavily_api_key: SecretStr | None = None) -> list[WebSearchResult]:
if tavily_api_key == None:
tavily_api_key = secret_str_from_env("TAVILY_API_KEY")

tavily = TavilyClient(api_key=tavily_api_key.get_secret_value() if tavily_api_key else None)
response = tavily.search(
query=query,
search_depth="advanced",
max_results=max_results,
include_raw_content=True,
)
response = _web_search(query=query, max_results=max_results, tavily_api_key=tavily_api_key)

transformed_results = [
WebSearchResult(
Expand All @@ -35,3 +25,21 @@ def web_search(query: str, max_results: int = 5, tavily_api_key: SecretStr | Non
]

return transformed_results


@tenacity.retry(stop=tenacity.stop_after_attempt(3), wait=tenacity.wait_fixed(1), reraise=True)
@persistent_inmemory_cache
def _web_search(query: str, max_results: int = 5, tavily_api_key: SecretStr | None = None) -> dict[str, t.Any]:
if tavily_api_key == None:
tavily_api_key = secret_str_from_env("TAVILY_API_KEY")

tavily = TavilyClient(api_key=tavily_api_key.get_secret_value() if tavily_api_key else None)
response: dict[str, t.Any] = tavily.search(
query=query,
search_depth="advanced",
max_results=max_results,
include_raw_content=True,
)

return response

Loading