Open WebUI with Qwen QwQ on Tavily Search

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134

"""
title: Tavily Search Tool
author: victor1203
maintainer: sievelau
description: This tool performs internet searches using the Tavily API to get real-time information with advanced context and Q&A capabilities
required_open_webui_version: 0.4.0
requirements: tavily-python
version: 1.1.0
licence: MIT
"""

from pydantic import BaseModel, Field
from typing import Optional, Literal
from tavily import TavilyClient


class Tools:
    def __init__(self):
        """Initialize the Tool with valves."""
        self.valves = self.Valves()
        self._client = None

    class Valves(BaseModel):
        tavily_api_key: str = Field(
            "", description="Your Tavily API key (starts with 'tvly-')"
        )
        search_depth: str = Field(
            "basic", description="Search depth - basic or advanced"
        )
        include_answer: bool = Field(
            True, description="Include an AI-generated answer in the response"
        )
        max_results: int = Field(
            5, description="Maximum number of search results to return (1-10)"
        )

    def _get_client(self) -> TavilyClient:
        """Get or create Tavily client instance."""
        if self._client is None:
            self._client = TavilyClient(api_key=self.valves.tavily_api_key)
        return self._client

    async def search(self, query: str, __event_emitter__=None) -> str:
        """
        Perform a Tavily search and return the results.

        Args:
            query: The search query string
            search_type: Type of search to perform:
                - regular: Standard search with full results
                - context: Optimized for RAG applications
                - qa: Quick answer to a specific question

        Returns:
            A formatted string containing the search results
        """
        try:
            # Input validation
            if not self.valves.tavily_api_key:
                return "Error: Tavily API key not configured. Please set up the API key in the tool settings."

            # Emit status that search is starting
            if __event_emitter__:
                await __event_emitter__(
                    {
                        "type": "status",
                        "data": {
                            "description": f"Initiating Tavily search...",
                            "done": False,
                        },
                    }
                )

            client = self._get_client()

            # Perform the search based on type
            if __event_emitter__:
                await __event_emitter__(
                    {
                        "type": "status",
                        "data": {
                            "description": "Fetching search results...",
                            "done": False,
                        },
                    }
                )
            result = client.search(
                query=query,
                search_depth=self.valves.search_depth,
                include_answer=self.valves.include_answer,
                max_results=self.valves.max_results,
            )

            # Format regular search results
            formatted_results = ""

            # Include AI answer if available
            if self.valves.include_answer and "answer" in result:
                formatted_results += f"AI Answer:\n{result['answer']}\n\n"

            # Add search results
            formatted_results += "Tavily Search Results:\n\n"
            for i, item in enumerate(
                result.get("results", [])[: self.valves.max_results], 1
            ):
                formatted_results += f"{i}. {item.get('title', 'No title')}\n"
                formatted_results += f"URL: {item.get('url', 'No URL')}\n"
                formatted_results += f"Content: {item.get('content', 'No content')}\n"
                formatted_results += f"Relevance Score: {item.get('score', '0.0')}\n"
                formatted_results += "\n"

            if __event_emitter__:
                await __event_emitter__(
                    {
                        "type": "status",
                        "data": {
                            "description": "Search completed successfully",
                            "done": True,
                        },
                    }
                )

            return formatted_results

        except Exception as e:
            error_message = f"An error occurred while performing the search: {str(e)}"
            if __event_emitter__:
                await __event_emitter__(
                    {
                        "type": "status",
                        "data": {"description": error_message, "done": True},
                    }
                )
            return error_message