From 1fed842fcca221e02c03f411821d3d7f89491213 Mon Sep 17 00:00:00 2001
From: Debanjum Singh Solanky <debanjum@gmail.com>
Date: Mon, 30 Sep 2024 01:19:56 -0700
Subject: [PATCH] Localize online search results to user country when location
 available

Get country code to server chat api from i.p location check on clients.
Use country code to get country specific online search results via Serper.dev API
---
 src/interface/desktop/chat.html           |  3 +++
 src/interface/desktop/shortcut.html       |  6 ++++--
 src/interface/obsidian/src/chat_view.ts   |  3 +++
 src/interface/web/app/chat/page.tsx       |  3 ++-
 src/interface/web/app/common/utils.ts     |  1 +
 src/interface/web/app/share/chat/page.tsx |  1 +
 src/khoj/processor/tools/online_search.py | 10 +++++-----
 src/khoj/routers/api_chat.py              |  6 ++++--
 src/khoj/utils/rawconfig.py               |  1 +
 9 files changed, 24 insertions(+), 10 deletions(-)

diff --git a/src/interface/desktop/chat.html b/src/interface/desktop/chat.html
index 3df00efc..cd47dae5 100644
--- a/src/interface/desktop/chat.html
+++ b/src/interface/desktop/chat.html
@@ -60,6 +60,7 @@
         let region = null;
         let city = null;
         let countryName = null;
+        let countryCode = null;
         let timezone = null;
         let chatMessageState = {
             newResponseTextEl: null,
@@ -76,6 +77,7 @@
                 region = data.region;
                 city = data.city;
                 countryName = data.country_name;
+                countryCode = data.country_code;
                 timezone = data.timezone;
             })
             .catch(err => {
@@ -157,6 +159,7 @@
                 ...(!!city && { city: city }),
                 ...(!!region && { region: region }),
                 ...(!!countryName && { country: countryName }),
+                ...(!!countryCode && { country_code: countryCode }),
                 ...(!!timezone && { timezone: timezone }),
             };
 
diff --git a/src/interface/desktop/shortcut.html b/src/interface/desktop/shortcut.html
index 86e5d906..4b07b8ea 100644
--- a/src/interface/desktop/shortcut.html
+++ b/src/interface/desktop/shortcut.html
@@ -308,18 +308,19 @@
     <script src="./utils.js"></script>
     <script src="./chatutils.js"></script>
     <script>
-
         let region = null;
         let city = null;
         let countryName = null;
+        let countryCode = null;
         let timezone = null;
 
         fetch("https://ipapi.co/json")
             .then(response => response.json())
             .then(data => {
-                region = data.region;
                 city = data.city;
+                region = data.region;
                 countryName = data.country_name;
+                countryCode = data.country_code;
                 timezone = data.timezone;
             })
             .catch(err => {
@@ -410,6 +411,7 @@
                 ...(!!city && { city: city }),
                 ...(!!region && { region: region }),
                 ...(!!countryName && { country: countryName }),
+                ...(!!countryCode && { country_code: countryCode }),
                 ...(!!timezone && { timezone: timezone }),
             };
 
diff --git a/src/interface/obsidian/src/chat_view.ts b/src/interface/obsidian/src/chat_view.ts
index 237fccf9..f27098ca 100644
--- a/src/interface/obsidian/src/chat_view.ts
+++ b/src/interface/obsidian/src/chat_view.ts
@@ -36,6 +36,7 @@ interface Location {
     region: string;
     city: string;
     countryName: string;
+    countryCode: string;
     timezone: string;
 }
 
@@ -70,6 +71,7 @@ export class KhojChatView extends KhojPaneView {
                     region: data.region,
                     city: data.city,
                     countryName: data.country_name,
+                    countryCode: data.country_code,
                     timezone: data.timezone,
                 };
             })
@@ -1060,6 +1062,7 @@ export class KhojChatView extends KhojPaneView {
                 city: this.location.city,
                 region: this.location.region,
                 country: this.location.countryName,
+                country_code: this.location.countryCode,
                 timezone: this.location.timezone,
             }),
         };
diff --git a/src/interface/web/app/chat/page.tsx b/src/interface/web/app/chat/page.tsx
index 0879956a..37a894e9 100644
--- a/src/interface/web/app/chat/page.tsx
+++ b/src/interface/web/app/chat/page.tsx
@@ -241,9 +241,10 @@ export default function Chat() {
             conversation_id: conversationId,
             stream: true,
             ...(locationData && {
+                city: locationData.city,
                 region: locationData.region,
                 country: locationData.country,
-                city: locationData.city,
+                country_code: locationData.countryCode,
                 timezone: locationData.timezone,
             }),
             ...(image64 && { image: image64 }),
diff --git a/src/interface/web/app/common/utils.ts b/src/interface/web/app/common/utils.ts
index c3395dbb..b5dbfec8 100644
--- a/src/interface/web/app/common/utils.ts
+++ b/src/interface/web/app/common/utils.ts
@@ -6,6 +6,7 @@ export interface LocationData {
     city: string;
     region: string;
     country: string;
+    countryCode: string;
     postal: string;
     latitude: number;
     longitude: number;
diff --git a/src/interface/web/app/share/chat/page.tsx b/src/interface/web/app/share/chat/page.tsx
index d9a945c5..116a7f7c 100644
--- a/src/interface/web/app/share/chat/page.tsx
+++ b/src/interface/web/app/share/chat/page.tsx
@@ -231,6 +231,7 @@ export default function SharedChat() {
                 region: locationData.region,
                 country: locationData.country,
                 city: locationData.city,
+                country_code: locationData.countryCode,
                 timezone: locationData.timezone,
             }),
             ...(image64 && { image: image64 }),
diff --git a/src/khoj/processor/tools/online_search.py b/src/khoj/processor/tools/online_search.py
index c7f7177f..3627c58e 100644
--- a/src/khoj/processor/tools/online_search.py
+++ b/src/khoj/processor/tools/online_search.py
@@ -7,7 +7,6 @@ from collections import defaultdict
 from typing import Callable, Dict, List, Optional, Tuple, Union
 
 import aiohttp
-import requests
 from bs4 import BeautifulSoup
 from markdownify import markdownify
 
@@ -80,7 +79,7 @@ async def search_online(
 
     with timer(f"Internet searches for {list(subqueries)} took", logger):
         search_func = search_with_google if SERPER_DEV_API_KEY else search_with_jina
-        search_tasks = [search_func(subquery) for subquery in subqueries]
+        search_tasks = [search_func(subquery, location) for subquery in subqueries]
         search_results = await asyncio.gather(*search_tasks)
         response_dict = {subquery: search_result for subquery, search_result in search_results}
 
@@ -115,8 +114,9 @@ async def search_online(
     yield response_dict
 
 
-async def search_with_google(query: str) -> Tuple[str, Dict[str, List[Dict]]]:
-    payload = json.dumps({"q": query})
+async def search_with_google(query: str, location: LocationData) -> Tuple[str, Dict[str, List[Dict]]]:
+    country_code = location.country_code.lower() if location and location.country_code else "us"
+    payload = json.dumps({"q": query, "gl": country_code})
     headers = {"X-API-KEY": SERPER_DEV_API_KEY, "Content-Type": "application/json"}
 
     async with aiohttp.ClientSession() as session:
@@ -220,7 +220,7 @@ async def read_webpage_with_jina(web_url: str) -> str:
             return response_json["data"]["content"]
 
 
-async def search_with_jina(query: str) -> Tuple[str, Dict[str, List[Dict]]]:
+async def search_with_jina(query: str, location: LocationData) -> Tuple[str, Dict[str, List[Dict]]]:
     encoded_query = urllib.parse.quote(query)
     jina_search_api_url = f"{JINA_SEARCH_API_URL}/{encoded_query}"
     headers = {"Accept": "application/json"}
diff --git a/src/khoj/routers/api_chat.py b/src/khoj/routers/api_chat.py
index a11ae54e..4ce50cf2 100644
--- a/src/khoj/routers/api_chat.py
+++ b/src/khoj/routers/api_chat.py
@@ -529,6 +529,7 @@ class ChatRequestBody(BaseModel):
     city: Optional[str] = None
     region: Optional[str] = None
     country: Optional[str] = None
+    country_code: Optional[str] = None
     timezone: Optional[str] = None
     image: Optional[str] = None
 
@@ -556,6 +557,7 @@ async def chat(
     city = body.city
     region = body.region
     country = body.country
+    country_code = body.country_code
     timezone = body.timezone
     image = body.image
 
@@ -653,8 +655,8 @@ async def chat(
 
         user_name = await aget_user_name(user)
         location = None
-        if city or region or country:
-            location = LocationData(city=city, region=region, country=country)
+        if city or region or country or country_code:
+            location = LocationData(city=city, region=region, country=country, country_code=country_code)
 
         if is_query_empty(q):
             async for result in send_llm_response("Please ask your query to get started."):
diff --git a/src/khoj/utils/rawconfig.py b/src/khoj/utils/rawconfig.py
index f0cd3962..15f5ea01 100644
--- a/src/khoj/utils/rawconfig.py
+++ b/src/khoj/utils/rawconfig.py
@@ -25,6 +25,7 @@ class LocationData(BaseModel):
     city: Optional[str]
     region: Optional[str]
     country: Optional[str]
+    country_code: Optional[str]
 
     def __str__(self):
         parts = []