From ad6f6bb0ed206c0f115761d0ee21020a867395cc Mon Sep 17 00:00:00 2001
From: Debanjum Singh Solanky <debanjum@gmail.com>
Date: Wed, 13 Mar 2024 15:30:24 +0530
Subject: [PATCH] Support webpage command in chat API

- Fallback to use webpage when SERPER not setup and online command was
  attempted
- Do not stop responding if can't retrieve online results. Try to
  respond without the online context
---
 src/khoj/processor/conversation/openai/gpt.py |  6 ++--
 src/khoj/routers/api_chat.py                  | 28 ++++++++++++++-----
 src/khoj/routers/helpers.py                   |  2 +-
 3 files changed, 25 insertions(+), 11 deletions(-)

diff --git a/src/khoj/processor/conversation/openai/gpt.py b/src/khoj/processor/conversation/openai/gpt.py
index dd01efd4..c3edf334 100644
--- a/src/khoj/processor/conversation/openai/gpt.py
+++ b/src/khoj/processor/conversation/openai/gpt.py
@@ -1,7 +1,7 @@
 import json
 import logging
 from datetime import datetime, timedelta
-from typing import Optional
+from typing import Dict, Optional
 
 from langchain.schema import ChatMessage
 
@@ -103,7 +103,7 @@ def send_message_to_model(messages, api_key, model, response_type="text"):
 def converse(
     references,
     user_query,
-    online_results: Optional[dict] = None,
+    online_results: Optional[Dict[str, Dict]] = None,
     conversation_log={},
     model: str = "gpt-3.5-turbo",
     api_key: Optional[str] = None,
@@ -141,7 +141,7 @@ def converse(
         completion_func(chat_response=prompts.no_online_results_found.format())
         return iter([prompts.no_online_results_found.format()])
 
-    if ConversationCommand.Online in conversation_commands:
+    if ConversationCommand.Online in conversation_commands or ConversationCommand.Webpage in conversation_commands:
         conversation_primer = (
             f"{prompts.online_search_conversation.format(online_results=str(online_results))}\n{conversation_primer}"
         )
diff --git a/src/khoj/routers/api_chat.py b/src/khoj/routers/api_chat.py
index 7a99869c..77888ce5 100644
--- a/src/khoj/routers/api_chat.py
+++ b/src/khoj/routers/api_chat.py
@@ -14,7 +14,11 @@ from khoj.database.adapters import ConversationAdapters, EntryAdapters, aget_use
 from khoj.database.models import KhojUser
 from khoj.processor.conversation.prompts import help_message, no_entries_found
 from khoj.processor.conversation.utils import save_to_conversation_log
-from khoj.processor.tools.online_search import search_online
+from khoj.processor.tools.online_search import (
+    online_search_enabled,
+    read_webpages,
+    search_online,
+)
 from khoj.routers.api import extract_references_and_questions
 from khoj.routers.helpers import (
     ApiUserRateLimiter,
@@ -274,7 +278,7 @@ async def chat(
     compiled_references, inferred_queries, defiltered_query = await extract_references_and_questions(
         request, common, meta_log, q, (n or 5), (d or math.inf), conversation_commands, location
     )
-    online_results: Dict = dict()
+    online_results: Dict[str, Dict] = {}
 
     if conversation_commands == [ConversationCommand.Notes] and not await EntryAdapters.auser_has_entries(user):
         no_entries_found_format = no_entries_found.format()
@@ -288,13 +292,23 @@ async def chat(
         conversation_commands.remove(ConversationCommand.Notes)
 
     if ConversationCommand.Online in conversation_commands:
+        if not online_search_enabled():
+            conversation_commands.remove(ConversationCommand.Online)
+            # If online search is not enabled, try to read webpages directly
+            if ConversationCommand.Webpage not in conversation_commands:
+                conversation_commands.append(ConversationCommand.Webpage)
+        else:
+            try:
+                online_results = await search_online(defiltered_query, meta_log, location)
+            except ValueError as e:
+                logger.warning(f"Error searching online: {e}. Attempting to respond without online results")
+
+    if ConversationCommand.Webpage in conversation_commands:
         try:
-            online_results = await search_online(defiltered_query, meta_log, location)
+            online_results = await read_webpages(defiltered_query, meta_log, location)
         except ValueError as e:
-            return StreamingResponse(
-                iter(["Please set your SERPER_DEV_API_KEY to get started with online searches 🌐"]),
-                media_type="text/event-stream",
-                status_code=200,
+            logger.warning(
+                f"Error directly reading webpages: {e}. Attempting to respond without online results", exc_info=True
             )
 
     if ConversationCommand.Image in conversation_commands:
diff --git a/src/khoj/routers/helpers.py b/src/khoj/routers/helpers.py
index a9dd5fb3..fb2a5df5 100644
--- a/src/khoj/routers/helpers.py
+++ b/src/khoj/routers/helpers.py
@@ -395,7 +395,7 @@ def generate_chat_response(
     q: str,
     meta_log: dict,
     compiled_references: List[str] = [],
-    online_results: Dict[str, Any] = {},
+    online_results: Dict[str, Dict] = {},
     inferred_queries: List[str] = [],
     conversation_commands: List[ConversationCommand] = [ConversationCommand.Default],
     user: KhojUser = None,