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
|
import json
from typing import List
import aiohttp
from pydantic import Field
from ...core.context import ContextProvider
from ...core.main import ContextItem, ContextItemDescription, ContextItemId
from .util import remove_meilisearch_disallowed_chars
class GoogleContextProvider(ContextProvider):
"""Type '@google' to reference the results of a Google search. For example, type "@google python tutorial" if you want to search and discuss ways of learning Python."""
title = "google"
display_title = "Google"
description = "Search Google"
dynamic = True
requires_query = True
serper_api_key: str = Field(
...,
description="Your SerpAPI key, used to programmatically make Google searches. You can get a key at https://serper.dev.",
)
_GOOGLE_CONTEXT_ITEM_ID = "google_search"
@property
def BASE_CONTEXT_ITEM(self):
return ContextItem(
content="",
description=ContextItemDescription(
name="Google Search",
description="Enter a query to search google",
id=ContextItemId(
provider_title=self.title, item_id=self._GOOGLE_CONTEXT_ITEM_ID
),
),
)
async def _google_search(self, query: str) -> str:
url = "https://google.serper.dev/search"
payload = json.dumps({"q": query})
headers = {"X-API-KEY": self.serper_api_key, "Content-Type": "application/json"}
async with aiohttp.ClientSession() as session:
async with session.post(url, headers=headers, data=payload) as response:
return await response.text()
async def provide_context_items(self, workspace_dir: str) -> List[ContextItem]:
return [self.BASE_CONTEXT_ITEM]
async def get_item(self, id: ContextItemId, query: str) -> ContextItem:
if not id.provider_title == self.title:
raise Exception("Invalid provider title for item")
results = await self._google_search(query)
json_results = json.loads(results)
content = f"Google Search: {query}\n\n"
if answerBox := json_results.get("answerBox"):
content += f"Answer Box ({answerBox['title']}): {answerBox['answer']}\n\n"
for result in json_results["organic"]:
content += f"{result['title']}\n{result['link']}\n{result['snippet']}\n\n"
ctx_item = self.BASE_CONTEXT_ITEM.copy()
ctx_item.content = content
ctx_item.description.id.item_id = remove_meilisearch_disallowed_chars(query)
return ctx_item
|
