Skip to content

Commit

Permalink
Chat improving (parse response message without json.loads, add `timeo…
Browse files Browse the repository at this point in the history
…ut` 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)
  • Loading branch information
deedy5 authored Jul 5, 2024
1 parent 744322d commit 8f2cd06
Show file tree
Hide file tree
Showing 3 changed files with 13 additions and 11 deletions.
3 changes: 2 additions & 1 deletion README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
5 changes: 3 additions & 2 deletions duckduckgo_search/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand All @@ -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
Expand All @@ -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}
Expand Down
16 changes: 8 additions & 8 deletions duckduckgo_search/duckduckgo_search.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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

Expand Down

0 comments on commit 8f2cd06

Please sign in to comment.