using google search api

prompt_engeneering.py

@@ -1,8 +1,9 @@
 from dataclasses import dataclass
 from datetime import datetime
-from typing import Generator
+from typing import Generator, Dict, List
 
-import requests
+from googleapiclient.discovery import build
+from streamlit import secrets
 
 INSTRUCTIONS = "Instructions: " \
                "Using the provided web search results, " \
@@ -12,30 +13,47 @@ INSTRUCTIONS = "Instructions: " \
                "write separate answers for each subject."
 
 
+def get_google_api_key():
+    """Returns the Google API key from streamlit's secrets"""
+    return secrets["google_search_api_key"]
+
+
+def get_google_cse_id():
+    """Returns the Google CSE ID from streamlit's secrets"""
+    return secrets["google_cse_id"]
+
+
+def google_search(search_term, **kwargs) -> list:
+    service = build("customsearch", "v1", developerKey=get_google_api_key())
+    search_engine = service.cse()
+    res = search_engine.list(q=search_term, cx=get_google_cse_id(), **kwargs).execute()
+    return res['items']
+
+
 @dataclass
 class SearchResult:
-    def __init__(self, title: str, body: str, url: str):
-        self.title = title
-        self.body = body
-        self.url = url
+    __slots__ = ["title", "body", "url"]
+    title: str
+    body: str
+    url: str
 
 
 def get_web_search_results(
-        prompt: str,
+        query: str,
         num_results: int,
 ) -> Generator[SearchResult, None, None]:
-    """Adds web search results to the prompt.
-    Using """
-    url = f"https://ddg-webapp-aagd.vercel.app/search?max_results={num_results}&q=${prompt}"
-    response = requests.get(url)
-    if response.status_code != 200:
-        raise ValueError(f"Failed to get web search results for prompt: {prompt}")
-    results = response.json()
-    for result in results:
+    """Gets a list of web search results using the Google search API"""
+    rew_results: List[Dict[str, str]] = google_search(
+        search_term=query,
+        num=num_results
+    )[:num_results]
+    for result in rew_results:
+        if result["snippet"].endswith("\xa0..."):
+            result["snippet"] = result["snippet"][:-4]
         yield SearchResult(
             title=result["title"],
-            body=result["body"],
-            url=result["href"],
+            body=result["snippet"],
+            url=result["link"],
         )
 
 
@@ -52,10 +70,10 @@ def rewrite_prompt(
 ) -> str:
     """Rewrites the prompt by adding web search results to it."""
     raw_results = get_web_search_results(
-        prompt=prompt,
+        query=prompt,
         num_results=5,
     )
-    formatted_results = "Web search results: " + format_search_result(raw_results)
+    formatted_results = "Web search results:\n" + format_search_result(raw_results)
     formatted_date = "Current date: " + datetime.now().strftime("%d/%m/%Y")
     formatted_prompt = f"Query: {prompt}"
     return "\n".join([formatted_results, formatted_date, INSTRUCTIONS, formatted_prompt])

requirements.txt

@@ -4,9 +4,10 @@ torch>1.12.1
 transformers~=4.26.0
 hatchling
 beautifulsoup4~=4.11.2
-urllib3
+urllib3~=1.26.14
 requests~=2.28.2
 accelerate
 bitsandbytes
-pytest
-sentencepiece
+pytest~=7.1.2
+sentencepiece
+google-api-python-client