Add prompt, tool to include nearby locations in google search query

Author: chatrtham

src/fastapi_app/llm_tools.py

@@ -11,7 +11,7 @@ def build_google_search_function() -> list[ChatCompletionToolParam]:
         {
             "type": "function",
             "function": {
-                "name": "search_database",
+                "name": "search_google",
                 "description": "Search for relevant products based on user query",
                 "parameters": {
                     "type": "object",
@@ -20,6 +20,17 @@ def build_google_search_function() -> list[ChatCompletionToolParam]:
                             "type": "string",
                             "description": "Query string to use for full text search, e.g. 'ตรวจสุขภาพ'",
                         },
+                        "locations": {
+                            "type": "array",
+                            "items": {
+                                "type": "string",
+                            },
+                            "description": """
+                            A list of nearby districts(Amphoes) from what the user provides.
+                            For example, if the user says `รังสิต`, the locations should be 
+                            [`ธัญบุรี`, `เมืองปทุมธานี`, `คลองหลวง`, `ลำลูกกา`].
+                            """
+                        }
                     },
                     "required": ["search_query"],
                 },
@@ -30,19 +41,19 @@ def build_google_search_function() -> list[ChatCompletionToolParam]:
 def extract_search_arguments(chat_completion: ChatCompletion):
     response_message = chat_completion.choices[0].message
     search_query = None
+
 
     if response_message.tool_calls:
         for tool in response_message.tool_calls:
             if tool.type != "function":
                 continue
             function = tool.function
-            if function.name == "search_database":
+            if function.name == "search_google":
                 arg = json.loads(function.arguments)
                 search_query = arg.get("search_query")
-    elif query_text := response_message.content:
-        search_query = query_text.strip()
+                locations = arg.get("locations", [])
 
-    return search_query
+    return search_query, locations
 
 
 def build_specify_package_function() -> list[ChatCompletionToolParam]:

src/fastapi_app/rag_advanced.py

@@ -65,17 +65,21 @@ async def google_search(self, messages):
             max_tokens=query_response_token_limit,
             n=1,
             tools=build_google_search_function(),
-            tool_choice="auto",
+            tool_choice={"type": "function", "function": {"name": "search_google"}},
         )
 
-        query_text = extract_search_arguments(query_chat_completion)
+        search_query, locations = extract_search_arguments(query_chat_completion)
+        locations = [f'"{location}"' for location in locations] if locations else []
+
+        query_text = f"{search_query} {' OR '.join(locations)}" if locations else search_query
 
         results = await self.searcher.google_search(query_text, top=3)
 
         sources_content = [f"[{(package.url)}]:{package.to_str_for_broad_rag()}\n\n" for package in results]
 
         thought_steps = [
-            ThoughtStep(title="Prompt to generate search arguments", description=query_text, props={}),
+            ThoughtStep(title="Prompt to generate search arguments", description=messages, props={}),
+            ThoughtStep(title="Google Search query", description=query_text, props={}),
             ThoughtStep(title="Google Search results", description=[result.to_dict() for result in results], props={}),
         ]
         return sources_content, thought_steps