Skip to content

Commit

Permalink
Add Baidu Search Tool for Phidata (#1468)
Browse files Browse the repository at this point in the history
- Implemented the `BaiduSearch` toolkit under `phi/tools/baidusearch.py`:
    - Allows to perform Baidu searches with options for maximum results, language settings, custom headers, and proxy support.
    - Integrated error handling for missing dependencies (`baidusearch`, `pycountry`) with installation instructions.
    - Configured search results to include title, URL, abstract, and ranking in JSON format.
    - Supports automatic language code conversion using `pycountry` (fallback to Chinese for invalid input).

- Added an example in `cookbook/tools/baidusearch_tools.py`:
    - Demonstrates how to use the `BaiduSearch` tool in a Phidata `Agent`.
    - Configured an agent to search in both English and Chinese and select the top 3 results from 5 searches.
    - Includes user-friendly prompts and markdown formatting for responses.

- Features:
    - Debug logging for easier troubleshooting.
    - Timeout settings for search requests.

This commit expands Phidata's search capabilities by integrating Baidu, making it useful for users in regions where Baidu is a preferred search engine.
  • Loading branch information
Vist233 authored Nov 27, 2024
1 parent db7e18d commit b36e08a
Show file tree
Hide file tree
Showing 2 changed files with 96 additions and 0 deletions.
14 changes: 14 additions & 0 deletions cookbook/tools/baidusearch_tools.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,14 @@
from phi.agent import Agent
from phi.tools.baidusearch import BaiduSearch

agent = Agent(
tools=[BaiduSearch()],
description="You are a search agent that helps users find the most relevant information using Baidu.",
instructions=[
"Given a topic by the user, respond with the 3 most relevant search results about that topic.",
"Search for 5 results and select the top 3 unique items.",
"Search in both English and Chinese.",
],
show_tool_calls=True,
)
agent.print_response("What are the latest advancements in AI?", markdown=True)
82 changes: 82 additions & 0 deletions phi/tools/baidusearch.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,82 @@
import json
from typing import Optional, List, Dict, Any

from phi.tools import Toolkit
from phi.utils.log import logger

try:
from baidusearch.baidusearch import search # type: ignore
except ImportError:
raise ImportError("`baidusearch` not installed. Please install using `pip install baidusearch`")

try:
from pycountry import pycountry
except ImportError:
raise ImportError("`pycountry` not installed. Please install using `pip install pycountry`")


class BaiduSearch(Toolkit):
"""
BaiduSearch is a toolkit for searching Baidu easily.
Args:
fixed_max_results (Optional[int]): A fixed number of maximum results.
fixed_language (Optional[str]): A fixed language for the search results.
headers (Optional[Any]): Headers to be used in the search request.
proxy (Optional[str]): Proxy to be used in the search request.
debug (Optional[bool]): Enable debug output.
"""

def __init__(
self,
fixed_max_results: Optional[int] = None,
fixed_language: Optional[str] = None,
headers: Optional[Any] = None,
proxy: Optional[str] = None,
timeout: Optional[int] = 10,
debug: Optional[bool] = False,
):
super().__init__(name="baidusearch")
self.fixed_max_results = fixed_max_results
self.fixed_language = fixed_language
self.headers = headers
self.proxy = proxy
self.timeout = timeout
self.debug = debug
self.register(self.baidu_search)

def baidu_search(self, query: str, max_results: int = 5, language: str = "zh") -> str:
"""Execute Baidu search and return results
Args:
query (str): Search keyword
max_results (int, optional): Maximum number of results to return, default 5
language (str, optional): Search language, default Chinese
Returns:
str: A JSON formatted string containing the search results.
"""
max_results = self.fixed_max_results or max_results
language = self.fixed_language or language

if len(language) != 2:
try:
language = pycountry.languages.lookup(language).alpha_2
except LookupError:
language = "zh"

logger.debug(f"Searching Baidu [{language}] for: {query}")

results = search(keyword=query, num_results=max_results)

res: List[Dict[str, str]] = []
for idx, item in enumerate(results, 1):
res.append(
{
"title": item.get("title", ""),
"url": item.get("url", ""),
"abstract": item.get("abstract", ""),
"rank": str(idx),
}
)
return json.dumps(res, indent=2)

0 comments on commit b36e08a

Please sign in to comment.