diff --git a/minimax_search_browse.py b/minimax_search_browse.py index b2cc988..799001e 100644 --- a/minimax_search_browse.py +++ b/minimax_search_browse.py @@ -9,6 +9,7 @@ import random import requests from openai import OpenAI +from tavily import TavilyClient from transformers import AutoTokenizer from concurrent.futures import ThreadPoolExecutor, as_completed @@ -52,6 +53,20 @@ def search_google(query): return response.json()["organic"] +def search_tavily(query): + """Search using Tavily API, returning results normalized to Serper shape.""" + tavily_client = TavilyClient(api_key=os.environ.get("TAVILY_API_KEY", "")) + response = tavily_client.search(query=query, max_results=10) + results = [] + for item in response.get("results", []): + results.append({ + "title": item.get("title", ""), + "url": item.get("url", ""), + "snippet": item.get("content", ""), + }) + return results + + def get_brief_text(contents): """Extract brief text from search results""" source_text = "" @@ -123,7 +138,11 @@ def get_search_results(query, max_retry=3): time.sleep(random.uniform(0, 16)) for retry_cnt in range(max_retry): try: - result = search_google(query) + search_provider = os.environ.get("SEARCH_PROVIDER", "").lower() + use_tavily = search_provider == "tavily" or ( + search_provider != "serper" and os.environ.get("TAVILY_API_KEY") + ) + result = search_tavily(query) if use_tavily else search_google(query) source_text = get_brief_text(result) break except Exception as e: diff --git a/pyproject.toml b/pyproject.toml index cc01259..89eea85 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -11,6 +11,7 @@ dependencies = [ "pyyaml>=6.0", "openai>=1.0.0", "transformers>=4.30.0", + "tavily-python>=0.3.0", ] [project.scripts]