From 8f2cd065d17b31d829e134d232a55b7da624df5b Mon Sep 17 00:00:00 2001
From: deedy5 <65482418+deedy5@users.noreply.github.com>
Date: Fri, 5 Jul 2024 12:54:22 +0300
Subject: [PATCH] Chat improving (parse response message without json.loads,
 add `timeout` parameter) (#235)

* chat(): parse response message without json.loads
* chat(): add 'timeout' parameter
* CLI: chat(): add `--timeout` parameter
* README: update `chat` section (add `timeout` parameter)
---
 README.md                              |  3 ++-
 duckduckgo_search/cli.py               |  5 +++--
 duckduckgo_search/duckduckgo_search.py | 16 ++++++++--------
 3 files changed, 13 insertions(+), 11 deletions(-)

diff --git a/README.md b/README.md
index bd7d9c1..decf2a6 100755
--- a/README.md
+++ b/README.md
@@ -225,13 +225,14 @@ Exceptions:
 ## 1. chat() - AI chat
 
 ```python
-def chat(self, keywords: str, model: str = "gpt-3.5") -> str:
+def chat(self, keywords: str, model: str = "gpt-3.5", timeout: int = 20) -> str:
     """Initiates a chat session with DuckDuckGo AI.
 
     Args:
         keywords (str): The initial message or question to send to the AI.
         model (str): The model to use: "gpt-3.5", "claude-3-haiku", "llama-3-70b", "mixtral-8x7b".
             Defaults to "gpt-3.5".
+        timeout (int): Timeout value for the HTTP client. Defaults to 20.
 
     Returns:
         str: The response from the AI.
diff --git a/duckduckgo_search/cli.py b/duckduckgo_search/cli.py
index a3021a2..a146966 100644
--- a/duckduckgo_search/cli.py
+++ b/duckduckgo_search/cli.py
@@ -136,6 +136,7 @@ def version():
 @click.option("-l", "--load", is_flag=True, default=False, help="load the last conversation from the json cache")
 @click.option("-p", "--proxy", default=None, help="the proxy to send requests, example: socks5://127.0.0.1:9150")
 @click.option("-ml", "--multiline", is_flag=True, default=False, help="multi-line input")
+@click.option("-t", "--timeout", default=20, help="timeout value for the HTTP client")
 @click.option(
     "-m",
     "--model",
@@ -149,7 +150,7 @@ def version():
     show_choices=False,
     default="3",
 )
-def chat(load, proxy, multiline, model):
+def chat(load, proxy, multiline, timeout, model):
     """CLI function to perform an interactive AI chat using DuckDuckGo API."""
     cache_file = "ddgs_chat_conversation.json"
     proxy = "socks5://127.0.0.1:9150" if proxy == "tb" else proxy
@@ -173,7 +174,7 @@ def chat(load, proxy, multiline, model):
         else:
             user_input = input()
         if user_input.strip():
-            resp_answer = client.chat(keywords=user_input, model=model)
+            resp_answer = client.chat(keywords=user_input, model=model, timeout=timeout)
             click.secho(f"AI: {resp_answer}", bg="black", fg="green")
 
             cache = {"vqd": client._chat_vqd, "messages": client._chat_messages}
diff --git a/duckduckgo_search/duckduckgo_search.py b/duckduckgo_search/duckduckgo_search.py
index 8cd0dec..1a12d0b 100644
--- a/duckduckgo_search/duckduckgo_search.py
+++ b/duckduckgo_search/duckduckgo_search.py
@@ -120,13 +120,14 @@ def _get_vqd(self, keywords: str) -> str:
         resp_content = self._get_url("POST", "https://duckduckgo.com", data={"q": keywords})
         return _extract_vqd(resp_content, keywords)
 
-    def chat(self, keywords: str, model: str = "gpt-3.5") -> str:
+    def chat(self, keywords: str, model: str = "gpt-3.5", timeout: int = 20) -> str:
         """Initiates a chat session with DuckDuckGo AI.
 
         Args:
             keywords (str): The initial message or question to send to the AI.
             model (str): The model to use: "gpt-3.5", "claude-3-haiku", "llama-3-70b", "mixtral-8x7b".
                 Defaults to "gpt-3.5".
+            timeout (int): Timeout value for the HTTP client. Defaults to 20.
 
         Returns:
             str: The response from the AI.
@@ -149,16 +150,15 @@ def chat(self, keywords: str, model: str = "gpt-3.5") -> str:
             "messages": self._chat_messages,
         }
         resp = self.client.post(
-            "https://duckduckgo.com/duckchat/v1/chat", headers={"x-vqd-4": self._chat_vqd}, json=json_data
+            "https://duckduckgo.com/duckchat/v1/chat",
+            headers={"x-vqd-4": self._chat_vqd},
+            json=json_data,
+            timeout=timeout,
         )
         self._chat_vqd = resp.headers.get("x-vqd-4", "")
 
-        messages = [
-            json_loads(x).get("message", "")
-            for line in resp.text.replace("data: ", "").replace("[DONE]", "").split("\n\n")
-            if (x := line.strip())
-        ]
-        result = "".join(messages)
+        messages = [e.split('","')[0] for e in resp.text.split('"message":"')[1:]]
+        result = "".join(messages).replace("\\n", "\n").replace("\\t", "\t")
         self._chat_messages.append({"role": "assistant", "content": result})
         return result