search.py

from tavily import TavilyClient
import ui
from typing import List, Dict
import os
from dotenv import load_dotenv
load_dotenv()

tavily_search = TavilyClient(api_key=os.getenv("TAVILY_API_KEY"))


def format_search_results(results: List[Dict]) -> str:
    formatted_results = []
    for result in results:
        formatted_result = ''
        formatted_result = f"[{result.get('number', '')}] Title: {result.get('title', 'N/A')}\n"
        formatted_result += f"   Snippet: {result.get('body', '')}\n"
        formatted_result += f"   URL: {result.get('href', 'N/A')}\n\n"
        formatted_results.append(formatted_result)
    return "\n".join(formatted_results)
    
def fetch_search_results(query: List[str], task_status, task_name, fn):
   
    results = []
    for q in query:
        task_status[task_name]["name"] = "Searching for: " + q
        yield from fn()
        ui.system_sub_update(f"Searching for: {q}")
        try:
            result = tavily_search.search(q)
            result = result["results"][:3]
            result = [{'title': res['title'], 'body': res['content'], 'href': res['url']} for res in result]
            results.extend(result)
        except:
            result = tavily_search.search(q)
            result = result["results"][:3]
            result = [{'title': res['title'], 'body': res['content'], 'href': res['url']} for res in result]
            results.extend(result)
    results = [{'number' : i+1, **result } for i, result in enumerate(results)]
    formatted_results = format_search_results(results)

    return formatted_results, results


if __name__ == "__main__":
    fetch_search_results(['gitlab', 'github'])