from openai import OpenAI
from smolagents import tool
import os
from dotenv import load_dotenv

load_dotenv()

PERPLEXITY_API_KEY = os.getenv("PERPLEXITY_API_KEY")


@tool
def perplexity_search_tool(query: str) -> str:
    """
    Performs a web search using the Perplexity API and returns the result.

    Args:
        query: The search query.

    Returns:
        The search result from Perplexity API or an error message.
    """
    if not PERPLEXITY_API_KEY:
        return "Error: Perplexity API key not set. Please set the PERPLEXITY_API_KEY environment variable."

    messages = [
        {
            "role": "system",
            "content": (
                "You are a helpful assistant"
            ),
        },
        {
            "role": "user",
            "content": query,
        },
    ]

    try:
        client = OpenAI(api_key=PERPLEXITY_API_KEY, base_url="https://api.perplexity.ai")

        response = client.chat.completions.create(
            model="sonar-pro",  # Using sonar-pro as per documentation
            messages=messages,
        )
        # Assuming the main content is in the first choice's message
        if response.choices and len(response.choices) > 0:
            return response.choices[0].message.content
        else:
            return "Error: No response choices received from Perplexity API."
    except Exception as e:
        return f"Error calling Perplexity API: {str(e)}"


if __name__ == "__main__":
    # Example usage:
    # Make sure to set the PERPLEXITY_API_KEY environment variable before running directly
    if PERPLEXITY_API_KEY:
        search_query = "What are the latest advancements in AI?"
        result = perplexity_search_tool(search_query)
        print(f"Search query: {search_query}")
        print(f"Result:\\n{result}")

        search_query_stars = "How many stars are there in our galaxy?"
        result_stars = perplexity_search_tool(search_query_stars)
        print(f"Search query: {search_query_stars}")
        print(f"Result:\\n{result_stars}")
    else:
        print("Please set the PERPLEXITY_API_KEY environment variable to test the tool.")