Implement AI chat function

README.md

@@ -11,14 +11,15 @@ Search for words, documents, images, videos, news, maps and text translation usi
 * [DDGS and AsyncDDGS classes](#ddgs-and-asyncddgs-classes)
 * [Proxy](#proxy)
 * [Exceptions](#exceptions)
-* [1. text() - text search](#1-text---text-search-by-duckduckgocom)
-* [2. answers() - instant answers](#2-answers---instant-answers-by-duckduckgocom)
-* [3. images() - image search](#3-images---image-search-by-duckduckgocom)
-* [4. videos() - video search](#4-videos---video-search-by-duckduckgocom)
-* [5. news() - news search](#5-news---news-search-by-duckduckgocom)
-* [6. maps() - map search](#6-maps---map-search-by-duckduckgocom)
-* [7. translate() - translation](#7-translate---translation-by-duckduckgocom)
-* [8. suggestions() - suggestions](#8-suggestions---suggestions-by-duckduckgocom)
+* [1. chat() - AI chat](#1-chat---ai-chat)
+* [2. text() - text search](#2-text---text-search-by-duckduckgocom)
+* [3. answers() - instant answers](#3-answers---instant-answers-by-duckduckgocom)
+* [4. images() - image search](#4-images---image-search-by-duckduckgocom)
+* [5. videos() - video search](#5-videos---video-search-by-duckduckgocom)
+* [6. news() - news search](#6-news---news-search-by-duckduckgocom)
+* [7. maps() - map search](#7-maps---map-search-by-duckduckgocom)
+* [8. translate() - translation](#8-translate---translation-by-duckduckgocom)
+* [9. suggestions() - suggestions](#9-suggestions---suggestions-by-duckduckgocom)
 
 ## Install
 ```python
@@ -44,6 +45,8 @@ python -m duckduckgo_search --help
 
 CLI examples:
 ```python3
+# AI chat
+ddgs chat
 # text search
 ddgs text -k "axon regeneration"
 # find and download pdf files via proxy (example: Tor browser)
@@ -223,7 +226,31 @@ Exceptions:
 
 [Go To TOP](#TOP)
 
-## 1. text() - text search by duckduckgo.com
+## 1. chat() - AI chat
+
+```python
+def chat(self, keywords: str, model: str = "gpt-3.5") -> 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". Defaults to "gpt-3.5".
+
+    Returns:
+        str: The response from the AI.
+    """
+```
+***Example***
+```python
+results = DDGS().chat("summarize Daniel Defoe's The Consolidator", model='claude-3-haiku')
+
+# async
+results = await AsyncDDGS().achat('describe the characteristic habits and behaviors of humans as a species')
+```
+
+[Go To TOP](#TOP)
+
+## 2. text() - text search by duckduckgo.com
 
 ```python
 def text(
@@ -263,7 +290,7 @@ results = await AsyncDDGS().atext('sun', region='wt-wt', safesearch='off', timel
 
 [Go To TOP](#TOP)
 
-## 2. answers() - instant answers by duckduckgo.com
+## 3. answers() - instant answers by duckduckgo.com
 
 ```python
 def answers(keywords: str) -> List[Dict[str, str]]:
@@ -286,7 +313,7 @@ results = await AsyncDDGS().aanswers("sun")
 
 [Go To TOP](#TOP)
 
-## 3. images() - image search by duckduckgo.com
+## 4. images() - image search by duckduckgo.com
 
 ```python
 def images(
@@ -344,7 +371,7 @@ results = await AsyncDDGS().aimages('sun', region='wt-wt', safesearch='off', max
 
 [Go To TOP](#TOP)
 
-## 4. videos() - video search by duckduckgo.com
+## 5. videos() - video search by duckduckgo.com
 
 ```python
 def videos(
@@ -391,7 +418,7 @@ results = await AsyncDDGS().avideos('sun', region='wt-wt', safesearch='off', tim
 
 [Go To TOP](#TOP)
 
-## 5. news() - news search by duckduckgo.com
+## 6. news() - news search by duckduckgo.com
 
 ```python
 def news(
@@ -424,7 +451,7 @@ results = await AsyncDDGS().anews('sun', region='wt-wt', safesearch='off', timel
 
 [Go To TOP](#TOP)
 
-## 6. maps() - map search by duckduckgo.com
+## 7. maps() - map search by duckduckgo.com
 
 ```python
 def maps(
@@ -472,7 +499,7 @@ results = await AsyncDDGS().amaps('shop', place="Baltimor", max_results=10)
 
 [Go To TOP](#TOP)
 
-## 7. translate() - translation by duckduckgo.com
+## 8. translate() - translation by duckduckgo.com
 
 ```python
 def translate(
@@ -505,7 +532,7 @@ results = await AsyncDDGS().atranslate('sun', to="de")
 
 [Go To TOP](#TOP)
 
-## 8. suggestions() - suggestions by duckduckgo.com
+## 9. suggestions() - suggestions by duckduckgo.com
 
 ```python
 def suggestions(

duckduckgo_search/cli.py

@@ -3,13 +3,14 @@
 import os
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from datetime import datetime
+from pathlib import Path
 from urllib.parse import unquote
 
 import click
 import pyreqwest_impersonate as pri
 
 from .duckduckgo_search import DDGS
-from .utils import json_dumps
+from .utils import json_dumps, json_loads
 from .version import __version__
 
 logger = logging.getLogger(__name__)
@@ -130,6 +131,45 @@ def version():
     return __version__
 
 
+@cli.command()
+@click.option("-s", "--save", is_flag=True, default=False, help="save the conversation in the json file")
+@click.option("-p", "--proxy", default=None, help="the proxy to send requests, example: socks5://localhost:9150")
+def chat(save, proxy):
+    """CLI function to perform an interactive AI chat using DuckDuckGo API."""
+    cache_file = "ddgs_chat_conversation.json"
+    models = ["gpt-3.5", "claude-3-haiku"]
+    client = DDGS(proxy=proxy)
+
+    print("DuckDuckGo AI chat. Available models:")
+    for idx, model in enumerate(models, start=1):
+        print(f"{idx}. {model}")
+    chosen_model_idx = input("Choose a model by entering its number[1]: ")
+    chosen_model_idx = 0 if not chosen_model_idx.strip() else int(chosen_model_idx) - 1
+    model = models[chosen_model_idx]
+    print(f"Using model: {model}")
+
+    if save and Path(cache_file).exists():
+        with open(cache_file) as f:
+            cache = json_loads(f.read())
+            client._chat_vqd = cache.get("vqd", None)
+            client._chat_messages = cache.get("messages", [])
+
+    while True:
+        user_input = input(f"{'-'*78}\nYou: ")
+        if not user_input.strip():
+            break
+
+        resp_answer = client.chat(keywords=user_input, model=model)
+        text = click.wrap_text(resp_answer, width=78, preserve_paragraphs=True)
+        click.secho(f"AI: {text}", bg="black", fg="green", overline=True)
+
+        cache = {"vqd": client._chat_vqd, "messages": client._chat_messages}
+        _save_json(cache_file, cache)
+
+        if "exit" in user_input.lower() or "quit" in user_input.lower():
+            break
+
+
 @cli.command()
 @click.option("-k", "--keywords", required=True, help="text search, keywords for query")
 @click.option("-r", "--region", default="wt-wt", help="wt-wt, us-en, ru-ru, etc. -region https://duckduckgo.com/params")

duckduckgo_search/duckduckgo_search.py

@@ -71,6 +71,8 @@ def __init__(
             verify=False,
         )
         self._exception_event = Event()
+        self._chat_messages: List[Dict[str, str]] = []
+        self._chat_vqd: str = ""
 
     def __enter__(self) -> "DDGS":
         return self
@@ -118,6 +120,44 @@ 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:
+        """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". Defaults to "gpt-3.5".
+
+        Returns:
+            str: The response from the AI.
+        """
+        models = {"claude-3-haiku": "claude-3-haiku-20240307", "gpt-3.5": "gpt-3.5-turbo-0125"}
+        # vqd
+        if not self._chat_vqd:
+            resp = self.client.get("https://duckduckgo.com/duckchat/v1/status", headers={"x-vqd-accept": "1"})
+            self._chat_vqd = resp.headers.get("x-vqd-4", "")
+
+        self._chat_messages.append({"role": "user", "content": keywords})
+
+        json_data = {
+            "model": models[model],
+            "messages": self._chat_messages,
+        }
+        resp = self.client.post(
+            "https://duckduckgo.com/duckchat/v1/chat", headers={"x-vqd-4": self._chat_vqd}, json=json_data
+        )
+        self._chat_vqd = resp.headers.get("x-vqd-4", "")
+
+        messages = []
+        for line in resp.text.replace("data: ", "").replace("[DONE]", "").split("\n\n"):
+            x = line.strip()
+            if x:
+                j = json_loads(x)
+                message = j.get("message", "")
+                messages.append(message)
+        result = "".join(messages)
+        self._chat_messages.append({"role": "assistant", "content": result})
+        return result
+
     def text(
         self,
         keywords: str,

duckduckgo_search/duckduckgo_search_async.py

@@ -36,6 +36,19 @@ async def __aexit__(
     ) -> None:
         pass
 
+    async def achat(self, keywords: str, model: str = "gpt-3.5") -> str:
+        """Initiates async 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". Defaults to "gpt-3.5".
+
+        Returns:
+            str: The response from the AI.
+        """
+        result = await self._loop.run_in_executor(self._executor, super().chat, keywords, model)
+        return result
+
     async def atext(
         self,
         keywords: str,

tests/test_cli.py

@@ -14,13 +14,18 @@
 @pytest.fixture(autouse=True)
 def pause_between_tests():
     time.sleep(0.5)
-    
+
 
 def test_version_command():
     result = runner.invoke(cli, ["version"])
     assert result.output.strip() == __version__
 
 
+def test_chat_command():
+    result = runner.invoke(cli, ["chat"])
+    assert "chat" in result.output
+
+
 def test_text_command():
     result = runner.invoke(cli, ["text", "-k", "python"])
     assert "title" in result.output

tests/test_duckduckgo_search.py

@@ -14,6 +14,9 @@ def test_context_manager():
         results = ddgs.news("cars", max_results=30)
         assert 27 <= len(results) <= 30
 
+def test_chat():
+    results = DDGS().chat("cat")
+    assert  len(results) >= 1
 
 def test_text():
     results = DDGS().text("cat", safesearch="off", timelimit="m", max_results=30)

tests/test_duckduckgo_search_async.py

@@ -7,6 +7,11 @@
 def pause_between_tests():
     time.sleep(0.5)
 
+@pytest.fixture(autouse=True)
+async def test_chat():
+    results = await AsyncDDGS().chat("cat")
+    assert  len(results) >= 1
+
 @pytest.mark.asyncio
 async def test_text():
     results = await AsyncDDGS().atext("sky", safesearch="off", timelimit="m", max_results=30)