from dataclasses import dataclass from datetime import datetime from typing import Generator import requests INSTRUCTIONS = "Instructions: " \ "Using the provided web search results, " \ "write a comprehensive reply to the given query. " \ "Make sure to cite results using [[number](URL)] notation after the reference. " \ "If the provided search results refer to multiple subjects with the same name, " \ "write separate answers for each subject." @dataclass class SearchResult: def __init__(self, title: str, body: str, url: str): self.title = title self.body = body self.url = url def get_web_search_results( prompt: str, num_results: int, ) -> Generator[SearchResult, None, None]: """Adds web search results to the prompt. Using """ url = f"https://ddg-webapp-aagd.vercel.app/search?max_results={num_results}&q=${prompt}" response = requests.get(url) if response.status_code != 200: raise ValueError(f"Failed to get web search results for prompt: {prompt}") results = response.json() for result in results: yield SearchResult( title=result["title"], body=result["body"], url=result["href"], ) def format_search_result(search_result: Generator[SearchResult, None, None]) -> str: """Formats a search result to be added to the prompt.""" ans = "" for i, result in enumerate(search_result): ans += f"[{i}] {result.body}\nURL: {result.url}\n\n" return ans def rewrite_prompt( prompt: str, ) -> str: """Rewrites the prompt by adding web search results to it.""" raw_results = get_web_search_results( prompt=prompt, num_results=5, ) formatted_results = "Web search results: " + format_search_result(raw_results) formatted_date = "Current date: " + datetime.now().strftime("%d/%m/%Y") formatted_prompt = f"Query: {prompt}" return "\n".join([formatted_results, formatted_date, INSTRUCTIONS, formatted_prompt])