search.py

import requests
from pydantic import BaseModel, Field
from typing import Callable, Any, Optional


class Valves(BaseModel):
    SEARXNG_URL: str = Field(
        default="https://search.ricky.cloudns.be/search",
        description="SearXNG search URL",
    )
    BAIDU_URL: str = Field(
        default="https://www.baidu.com/s", description="Baidu search URL"
    )
    # 添加 Bing 搜索 URL
    BING_URL: str = Field(
        default="https://www.bing.com/search", description="Bing search URL"
    )
    TIMEOUT: int = Field(default=30, description="Request timeout in seconds")


class Tools:
    def __init__(self):
        self.valves = Valves()
        self.reader_api = "https://r.jina.ai/"
        self.citation = True

    async def read_url(
        self, url: str, __event_emitter__: Optional[Callable[[dict], Any]] = None
    ) -> str:
        """
        Read and extract the main content from a given URL.
        :param url: The URL to read from.
        :return: The main content of the page in clean, LLM-friendly text.
        """
        try:
            if __event_emitter__:
                await __event_emitter__(
                    {
                        "type": "status",
                        "data": {
                            "description": f"Reading content from {url}",
                            "status": "in_progress",
                        },
                    }
                )

            data = {"url": url}
            response = requests.post(
                self.reader_api, data=data, timeout=self.valves.TIMEOUT
            )
            response.raise_for_status()
            content = response.text

            if __event_emitter__:
                await __event_emitter__(
                    {
                        "type": "status",
                        "data": {
                            "description": "Content retrieved",
                            "status": "complete",
                        },
                    }
                )

            result_presentation = """
Result Presentation
STRICTLY ADHERE TO THIS FORMAT:
### [Site Name - Title](URL) 
#### 核心内容（Key Insights）:
- Point 1
- Point 2
...
#### 总结（Summary）:
---"""

            return f"{content}\n---\n{result_presentation}"
        except Exception as e:
            if __event_emitter__:
                await __event_emitter__(
                    {
                        "type": "status",
                        "data": {
                            "description": f"Error reading URL: {str(e)}",
                            "status": "error",
                        },
                    }
                )
            return f"Error reading URL: {str(e)}"

    async def search(
        self,
        search_term: str,
        engine: str = "bing",
        __event_emitter__: Optional[Callable[[dict], Any]] = None,
    ) -> str:
        """
        Issues a query to a search engine and displays the results. (Note: Use this when the query requires gathering information from external sources, such as current events, unfamiliar terms, or when explicit browsing is requested.)
        :param search_term: The search term or phrase to look up. (Note: think the keywords most likely to help users retrieve information as the "search_term",don't include search engine)
        :param engine: The search engine to use ('bing', 'baidu', or 'google'). Note:bing(Default, suitable for most queries), baidu(Best for content about China), google(For backup use only)
        :return: The search results in text format.
        """
        try:
            if __event_emitter__:
                await __event_emitter__(
                    {
                        "type": "status",
                        "data": {
                            "description": f"Searching with {engine}: {search_term}",
                            "status": "in_progress",
                        },
                    }
                )

            if engine == "baidu":
                url = f"{self.reader_api}{self.valves.BAIDU_URL}?wd={search_term}"
                headers = {"X-Target-Selector": "#content_left"}
            elif engine == "bing":
                # 添加 Bing 搜索的处理
                url = f"{self.reader_api}{self.valves.BING_URL}?q={search_term}"
                headers = {"X-Target-Selector": "#b_results"}
            else:
                url = f"{self.reader_api}{self.valves.SEARXNG_URL}?q={search_term}"
                headers = {"X-Target-Selector": "#urls"}

            response = requests.get(url, headers=headers, timeout=self.valves.TIMEOUT)
            response.raise_for_status()
            content = response.text

            if __event_emitter__:
                await __event_emitter__(
                    {
                        "type": "status",
                        "data": {
                            "description": f"Search completed with {engine}",
                            "status": "complete",
                        },
                    }
                )

            search_result_processing = """
Search Result Processing:
- Identify 1-3 highly relevant results from the search.
- Focus on the most informative and reliable sources.
- Integrate the key points from these sources to formulate a comprehensive answer to the query.
- If applicable, mention additional relevant information inferred from visible snippets in the search results.
"""

            return f"{content}\n---\n{search_result_processing}"

        except Exception as e:
            if __event_emitter__:
                await __event_emitter__(
                    {
                        "type": "status",
                        "data": {
                            "description": f"Error during search: {str(e)}",
                            "status": "error",
                        },
                    }
                )
            return f"Error during search: {str(e)}"