import modal import os from suno_utils.worker.modal_base import get_modal_base_image from openai import OpenAI # type: ignore[attr-defined] ############## CHANGE THESE ############## DEPLOYMENT_TYPE = "dev" # dev, prod APP_NAME = f"brave-search-{DEPLOYMENT_TYPE}" assert APP_NAME.endswith(DEPLOYMENT_TYPE) N_CPU = 1 aws_secret = modal.Secret.from_name("studio-aws") brave_secret = modal.Secret.from_name("brave-secret") tavily_secret = modal.Secret.from_name("tavily-search-secret") SECRETS = [ aws_secret, brave_secret, tavily_secret, modal.Secret.from_name("openai-secret"), modal.Secret.from_name("datadog-metrics"), modal.Secret.from_name("api-callback-token"), ] SEARCH_QUERY_SYSTEM_PROMPT = """ You are given a tweet and related contexts (maybe image and its description). Generate a concise web search query that will help someone understand the reference, meme, or context behind them. • If the tweet (or its image) explicitly mentions a date, include that exact date. • If there is no explicit date, you may prepend “recent” to capture recency, but do NOT invent or guess a month or year. Return only the search query — no labels, no explanations, no prefixes. Input: """ base_image = ( get_modal_base_image().pip_install("requests").add_local_python_source("suno_utils", copy=False) ) app = modal.App(APP_NAME, image=base_image, secrets=SECRETS) @app.cls( cpu=N_CPU, secrets=SECRETS, timeout=1000, scaledown_window=360, retries=modal.Retries( max_retries=2, backoff_coefficient=2.0, initial_delay=5.0, ), min_containers=1, ) @modal.concurrent(max_inputs=20) class SearchWorker: """Worker class for performing searches using different search engines.""" # API endpoint constants BRAVE_SEARCH_API_URL = "https://api.search.brave.com/res/v1/web/search" TAVILY_SEARCH_API_URL = "https://api.tavily.com/search" def __init__(self): self.openai_client = OpenAI(api_key=os.environ["OPENAI_API_KEY"]) @modal.method() def search_with_brave(self, query: str, count: int = 1) -> dict: """ Search the web using Brave Search API. Args: query: The search query string count: Number of results to return Returns: Dictionary containing search results """ import requests import re import urllib.parse # Clean and encode the query clean_query = re.sub(r"\s+", " ", query).strip() encoded_query = urllib.parse.quote(clean_query) url = self.BRAVE_SEARCH_API_URL headers = { "Accept": "application/json", "X-Subscription-Token": os.environ["BRAVE_API_KEY"], } params = { "q": encoded_query, "count": count, "search_lang": "en", } try: response = requests.get(url, headers=headers, params=params) response.raise_for_status() data = response.json() results = data.get("web", {}).get("results", []) if not results: return {"error": "No results found."} print(f"Brave search results:\n Query: {query}\n Results: {results}") # Return results return { "results": [ {"title": r.get("title"), "description": r.get("description")} for r in results ] } except requests.exceptions.RequestException as e: print(f"Brave search request failed: {str(e)}") return {"error": f"Search request failed: {str(e)}"} except (KeyError, ValueError) as e: print(f"Error parsing Brave search results: {str(e)}") return {"error": f"Error parsing search results: {str(e)}"} @modal.method() def search_with_tavily(self, query: str, max_results: int = 1, days: int = 7) -> dict: """ Search the web using Tavily API. Args: query: The search query string max_results: Maximum number of results to return days: Time range in days for search results Returns: Dictionary containing search results """ import requests import os url = self.TAVILY_SEARCH_API_URL payload = { "query": query, # Use the actual query parameter "topic": "general", "search_depth": "basic", "chunks_per_source": 3, "max_results": max_results, "time_range": None, "days": days, "include_answer": True, "include_raw_content": False, "include_images": False, "include_image_descriptions": False, "include_domains": [], "exclude_domains": [], } headers = { "Authorization": f"Bearer {os.environ.get('TAVILY_SEARCH_SECRET', '')}", "Content-Type": "application/json", } try: response = requests.post(url, json=payload, headers=headers) response.raise_for_status() data = response.json() results = data.get("results", []) if not results: # Check if the response has answer answer = data.get("answer", None) if answer: return {"results": [{"title": answer, "description": answer}]} return {"error": "No results found."} return { "results": [{"title": r.get("title"), "description": r.get("content")} for r in results] } except requests.exceptions.RequestException as e: print(f"Tavily search request failed: {str(e)}") return {"error": f"Tavily search request failed: {str(e)}"} except ValueError as e: print(f"Error parsing Tavily search results: {str(e)}") return {"error": f"Error parsing Tavily search results: {str(e)}"} @modal.method() def compose_search_query(self, context: str) -> str: """ Compose a search query based on the context. """ completion = self.openai_client.chat.completions.create( model="gpt-4o-mini", messages=[ {"role": "system", "content": SEARCH_QUERY_SYSTEM_PROMPT}, {"role": "user", "content": context}, ], ) return completion.choices[0].message.content @app.local_entrypoint() def main(): worker = SearchWorker() tweet = """ I know I am going to get a lot of emotional flack for this but I’m going to say it anyways. Trump did something deeply irresponsible. Crazy risk. Insane tails. But whether you like it or not we are now at war with China. """ result = worker.search_with_brave.remote(tweet, count=3) print(f"brave result: {result}") result = worker.search_with_tavily.remote(tweet, max_results=3, days=7) print(f"tavily result: {result}") context = """ Original Tweet Content: "o3 is AGI" https://t.co/2On739AeCU Tweet Image Description: **Image 1:** **Transcription:** How many 'r's are there in the word 'strawberry'? Thought for a second There are two "r"s in “strawberry.” """ result = worker.compose_search_query.remote(context) print(f"compose search query result: {result}") search_result = worker.search_with_tavily.remote(result, max_results=3, days=7) print(f"search result: {search_result}")