WIP- Under maintenence- Web research working.

lib/ai_web_researcher/tavily_ai_search.py

@@ -0,0 +1,156 @@
+"""
+This Python script uses the Tavily AI service to perform advanced searches based on specified keywords and options. It retrieves Tavily AI search results, pretty-prints them using Rich and Tabulate, and provides additional information such as the answer to the search query and follow-up questions.
+
+Features:
+- Utilizes the Tavily AI service for advanced searches.
+- Retrieves API keys from the environment variables loaded from a .env file.
+- Configures logging with Loguru for informative messages.
+- Implements a retry mechanism using Tenacity to handle transient failures during Tavily searches.
+- Displays search results, including titles, snippets, and links, in a visually appealing table using Tabulate and Rich.
+
+Usage:
+- Ensure the necessary API keys are set in the .env file.
+- Run the script to perform a Tavily AI search with specified keywords and options.
+- The search results, including titles, snippets, and links, are displayed in a formatted table.
+- Additional information, such as the answer to the search query and follow-up questions, is presented in separate tables.
+
+Modifications:
+- To modify the script, update the environment variables in the .env file with the required API keys.
+- Adjust the search parameters, such as keywords and search depth, in the `get_tavilyai_results` function as needed.
+- Customize logging configurations and table formatting according to preferences.
+
+To-Do (TBD):
+- Consider adding further enhancements or customization based on specific use cases.
+
+Note: This script depends on external libraries such as Tavily, Rich, Tabulate, Loguru, and Tenacity. Install them using 'pip install tavily rich tabulate loguru tenacity' if not already installed.
+"""
+
+
+import os
+from pathlib import Path
+import sys
+from dotenv import load_dotenv
+from loguru import logger
+from tavily import TavilyClient
+from rich import print
+from tabulate import tabulate
+# Load environment variables from .env file
+load_dotenv(Path('../../.env'))
+from rich import print
+
+# Configure logger
+logger.remove()
+logger.add(sys.stdout,
+           colorize=True,
+           format="<level>{level}</level>|<green>{file}:{line}:{function}</green>| {message}"
+           )
+from tenacity import retry, stop_after_attempt, wait_random_exponential
+
+from .gpt_titles_faq import gpt_titles_faqs_google_search
+
+@retry(wait=wait_random_exponential(min=1, max=60), stop=stop_after_attempt(6))
+def get_tavilyai_results(keywords, include_urls, search_depth="advanced"):
+    """
+    Get Tavily AI search results based on specified keywords and options.
+
+    Args:
+        keywords (str): Keywords for Tavily AI search.
+        include_urls (str): Comma-separated URLs to include in the search.
+        search_depth (str, optional): Search depth option (default is "advanced").
+
+    Returns:
+        dict: Tavily AI search results.
+    """
+    # Run Tavily search
+    logger.info(f"Running Tavily search on: {keywords}")
+
+    # Retrieve API keys
+    api_key = os.getenv('TAVILY_API_KEY')
+    if not api_key:
+        raise ValueError("API keys for Tavily or OpenAI are not set.")
+
+    # Initialize Tavily client
+    try:
+        client = TavilyClient(api_key=api_key)
+    except Exception as err:
+        logger.error(f"Failed to create Tavily client. Check TAVILY_API_KEY: {err}")
+        exit(1)
+    try:
+        if include_urls:
+            tavily_search_result = client.search(keywords, search_depth, include_answer=True, include_domains=include_urls)
+        else:
+            tavily_search_result = client.search(keywords, search_depth, include_answer=True)
+        print_result_table(tavily_search_result)
+        return(tavily_search_result)
+
+    except Exception as err:
+        logger.error(f"Failed to do Tavily Research: {err}")
+
+
+def print_result_table(output_data):
+    """ Pretty print the tavily AI serch result. """
+    # Prepare data for tabulate
+    table_data = []
+    for item in output_data.get("results"):
+        title = item.get("title", "")
+        snippet = item.get("content", "")
+        link = item.get("url", "")
+        table_data.append([title, snippet, link])
+    
+    # Define table headers
+    table_headers = ["Title", "Snippet", "Link"] 
+    # Display the table using tabulate
+    table = tabulate(table_data, 
+            headers=table_headers, 
+            tablefmt="fancy_grid", 
+            colalign=["left", "left", "left"], 
+            maxcolwidths=[30, 60, 30])
+    # Print the table
+    print(table)
+
+    # Save the combined table to a file
+    try:
+        save_in_file(table)
+    except Exception as save_results_err:
+        logger.error(f"Failed to save search results: {save_results_err}")
+    
+    # Display the 'answer' in a table
+    table_headers = [f"The answer to search query: {output_data.get('query')}"]
+    table_data = [[output_data.get("answer")]]
+    table = tabulate(table_data, 
+            headers=table_headers, 
+            tablefmt="fancy_grid", 
+            maxcolwidths=[80])
+    print(table)
+    # Save the combined table to a file
+    try:
+        save_in_file(table)
+    except Exception as save_results_err:
+        logger.error(f"Failed to save search results: {save_results_err}")
+    
+    # Display the 'follow_up_questions' in a table
+    table_headers = [f"Search Engine follow up questions for query: {output_data.get('query')}"]
+    table_data = [[output_data.get("follow_up_questions")]]
+    table = tabulate(table_data, 
+            headers=table_headers, 
+            tablefmt="fancy_grid",
+            maxcolwidths=[80])
+    print(table)
+    # Save the combined table to a file
+    try:
+        save_in_file(table)
+    except Exception as save_results_err:
+        logger.error(f"Failed to save search results: {save_results_err}")
+
+
+def save_in_file(table_content):
+    """ Helper function to save search analysis in a file. """
+    file_path = os.environ.get('SEARCH_SAVE_FILE')
+    try:
+        # Save the content to the file
+        with open(file_path, "a") as file:
+            file.write(table_content)
+            file.write("\n" * 3)  # Add three newlines at the end
+        logger.info(f"Search content saved to {file_path}")
+    except Exception as e:
+        logger.error(f"Error occurred while writing to the file: {e}")