utils_groupclassification/tavily_search_utils.py

import os

from dotenv import load_dotenv
from tavily import TavilyClient

load_dotenv(".env.dev")


# Function to perform a search
def tavily_search(query):
    tavily_client = TavilyClient(api_key=os.getenv("TAVILY_API_KEY"))
    search_result = tavily_client.get_search_context(
        query, search_depth="advanced", max_tokens=8000
    )
    return search_result