From 61cf88e16a1d70e301bcd7b88597976f1a26ba87 Mon Sep 17 00:00:00 2001 From: deedy5 <65482418+deedy5@users.noreply.github.com> Date: Wed, 15 May 2024 16:42:02 +0300 Subject: [PATCH] Implement AI `chat` function --- README.md | 59 ++++++++++++++------ duckduckgo_search/cli.py | 42 +++++++++++++- duckduckgo_search/duckduckgo_search.py | 40 +++++++++++++ duckduckgo_search/duckduckgo_search_async.py | 13 +++++ tests/test_cli.py | 7 ++- tests/test_duckduckgo_search.py | 3 + tests/test_duckduckgo_search_async.py | 5 ++ 7 files changed, 151 insertions(+), 18 deletions(-) diff --git a/README.md b/README.md index d3baf43..0f78fb0 100755 --- a/README.md +++ b/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( diff --git a/duckduckgo_search/cli.py b/duckduckgo_search/cli.py index 38d974a..deccd8c 100644 --- a/duckduckgo_search/cli.py +++ b/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") diff --git a/duckduckgo_search/duckduckgo_search.py b/duckduckgo_search/duckduckgo_search.py index 68b97d3..54ebc89 100644 --- a/duckduckgo_search/duckduckgo_search.py +++ b/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, diff --git a/duckduckgo_search/duckduckgo_search_async.py b/duckduckgo_search/duckduckgo_search_async.py index dd27a76..f558599 100644 --- a/duckduckgo_search/duckduckgo_search_async.py +++ b/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, diff --git a/tests/test_cli.py b/tests/test_cli.py index b240eb2..a50fa53 100644 --- a/tests/test_cli.py +++ b/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 diff --git a/tests/test_duckduckgo_search.py b/tests/test_duckduckgo_search.py index 8315387..dab4dc7 100644 --- a/tests/test_duckduckgo_search.py +++ b/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) diff --git a/tests/test_duckduckgo_search_async.py b/tests/test_duckduckgo_search_async.py index 19712ff..4286ad7 100644 --- a/tests/test_duckduckgo_search_async.py +++ b/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)