LangBot/plugins/misc/TavilySearch/components/tools/tavily_search.py

# Auto generated by LangBot Plugin SDK.
# Please refer to https://langbot.app/docs/en/plugin/dev/tutor.html for more details.
from __future__ import annotations

import asyncio
from typing import Any
from tavily import TavilyClient

from langbot_plugin.api.definition.components.tool.tool import Tool
from langbot_plugin.api.entities.builtin.provider import session as provider_session


class TavilySearchTool(Tool):
    """
    A tool for searching using Tavily API.
    Tavily is a search engine built specifically for AI agents (LLMs),
    delivering real-time, accurate, and factual results.
    """

    def _process_params(self, params: dict[str, Any]) -> dict[str, Any]:
        """
        Process and validate the search parameters.

        Args:
            params: The raw search parameters from LLM

        Returns:
            dict: Processed parameters ready for Tavily API
        """
        processed_params = {}

        # Required parameter
        if "query" in params:
            processed_params["query"] = params["query"]

        # Optional parameters with defaults
        processed_params["search_depth"] = params.get("search_depth", "basic")
        processed_params["topic"] = params.get("topic", "general")
        processed_params["max_results"] = params.get("max_results", 5)

        # Boolean flags
        if params.get("include_answer", False):
            processed_params["include_answer"] = True
        if params.get("include_images", False):
            processed_params["include_images"] = True
        if params.get("include_raw_content", False):
            processed_params["include_raw_content"] = True

        return processed_params

    def _format_results(self, search_results: dict, params: dict[str, Any]) -> str:
        """
        Format the search results into readable text.

        Args:
            search_results: Raw results from Tavily API
            params: Original search parameters

        Returns:
            str: Formatted search results as markdown
        """
        output_lines = []

        # Add answer if available
        if params.get("include_answer") and search_results.get("answer"):
            output_lines.append(f"**Answer:** {search_results['answer']}\n")

        # Add search results
        if "results" in search_results:
            for idx, result in enumerate(search_results["results"], 1):
                title = result.get("title", "No Title")
                url = result.get("url", "")
                content = result.get("content", "")
                score = result.get("score", "")

                output_lines.append(f"## Result {idx}: {title}\n")
                output_lines.append(f"**URL:** {url}\n")

                if score:
                    output_lines.append(f"**Relevance Score:** {score}\n")

                if content:
                    output_lines.append(f"**Content:**\n{content}\n")

                if params.get("include_raw_content") and result.get("raw_content"):
                    output_lines.append(f"**Raw Content:**\n{result['raw_content']}\n")

                output_lines.append("---\n")

        # Add images if requested
        if params.get("include_images") and search_results.get("images"):
            output_lines.append("**Images:**\n")
            for image in search_results["images"]:
                if isinstance(image, dict):
                    image_url = image.get("url")
                    description = image.get("description", "Tavily search result image")
                else:
                    image_url = image
                    description = "Tavily search result image"

                if image_url:
                    output_lines.append(f"- {image_url}\n")
            output_lines.append("\n")

        return "\n".join(output_lines)

    async def call(self, params: dict[str, Any], session: provider_session.Session, query_id: int) -> dict[str, Any] | str:
        """
        Execute the Tavily search.

        Args:
            params: Search parameters provided by LLM

        Returns:
            Search results as formatted string or error message
        """
        # Get API key from the current invocation's installation config.
        config = self.get_plugin_config()
        api_key = config.get("tavily_api_key")

        if not api_key:
            return "Error: Tavily API key is not configured. Please set it in the plugin configuration."

        # Validate required parameter
        query = params.get("query")
        if not query:
            return "Error: Please provide a search query."

        try:
            # Initialize Tavily client
            client = TavilyClient(api_key=api_key)

            # Process parameters
            processed_params = self._process_params(params)

            # Execute search off the event loop: TavilyClient.search is synchronous
            # and would otherwise stall every other tenant on this shared worker.
            search_results = await asyncio.to_thread(
                client.search, **processed_params
            )

            # Check if results exist
            if not search_results.get("results"):
                return f"No results found for '{query}' in Tavily."

            # Format and return results
            formatted_text = self._format_results(search_results, params)

            return formatted_text

        except Exception as e:
            return f"Error occurred while searching: {str(e)}"