From 423e557123eea018074cc8cc9155ce4cc832d8a8 Mon Sep 17 00:00:00 2001
From: LuoChen <luochen1990@gmail.com>
Date: Tue, 6 Aug 2024 23:25:45 +0800
Subject: [PATCH] add: [tools] google_search

---
 src/ai_powered/tools/google_search.py       | 43 +++++++++++++++++++++
 test/examples/chat_bot/use_google_search.py | 14 +++++++
 2 files changed, 57 insertions(+)
 create mode 100644 src/ai_powered/tools/google_search.py
 create mode 100644 test/examples/chat_bot/use_google_search.py

diff --git a/src/ai_powered/tools/google_search.py b/src/ai_powered/tools/google_search.py
new file mode 100644
index 0000000..ac58969
--- /dev/null
+++ b/src/ai_powered/tools/google_search.py
@@ -0,0 +1,43 @@
+import os
+from typing import List
+import requests
+from typing_extensions import Literal, TypedDict, NotRequired
+from ai_powered.tool_call import make_tool
+
+class SearchOption(TypedDict):
+    siteSearch: NotRequired[str]
+    dateRestrict: NotRequired[Literal["d[1]", "w[1]", "m[1]"]]
+
+class SearchResult(TypedDict):
+    title: str
+    link: str
+    snippet: str
+
+GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY", "mock_google_pse_api_key")
+GOOGLE_ENGINE_ID = os.getenv("GOOGLE_ENGINE_ID", "mock_google_pse_engine_id")
+
+@make_tool
+def google_search(keywords: str, options: SearchOption) -> List[SearchResult]:
+    ''' 用Google搜索互联网,对于专业问题请用英文关键词搜索 '''
+    url = "https://www.googleapis.com/customsearch/v1"
+    params : dict[str, str] = {
+        "key": GOOGLE_API_KEY,
+        "cx": GOOGLE_ENGINE_ID,
+        "q": keywords,
+        **options #type: ignore
+    }
+
+    response = requests.get(url, params=params)
+    response.raise_for_status()  # 如果请求失败，抛出异常
+
+    search_results = response.json()
+    results : list[SearchResult] = []
+
+    for item in search_results.get("items", []):
+        results.append(SearchResult(
+            title=item["title"],
+            link=item["link"],
+            snippet=item["snippet"]
+        ))
+
+    return results
diff --git a/test/examples/chat_bot/use_google_search.py b/test/examples/chat_bot/use_google_search.py
new file mode 100644
index 0000000..3111deb
--- /dev/null
+++ b/test/examples/chat_bot/use_google_search.py
@@ -0,0 +1,14 @@
+from ai_powered import ChatBot
+from ai_powered.colors import gray, green
+from ai_powered.tools.google_search import google_search
+
+class MyChatBot (ChatBot):
+    system_prompt = '''
+    Please answer the user's questions. If any calculations are required, use the calculator available in the tool. It supports complex Python expressions. When using it, make sure to convert the user's mathematical expression to a valid Python expression. Do not use any undefined functions; if the user's expression includes function calls, convert them to Python's built-in functions or syntax.
+    '''
+    tools = (google_search,)
+
+def test_use_google_search():
+    bot = MyChatBot()
+    print(green(bot.chat("what's USD price in CNY today?")))
+    print(gray(f"{bot.conversation}"))