danswer/backend/scripts/simulate_chat_frontend.py
2024-01-02 11:25:50 -08:00

93 lines
2.8 KiB
Python

# This file is purely for development use, not included in any builds
# Use this to test the chat feature
# This script does not allow for branching logic that is supported by the backend APIs
# This script also does not allow for editing/regeneration of user/model messages
# Have Danswer API server running to use this.
import json
import requests
from danswer.configs.app_configs import APP_PORT
LOCAL_CHAT_ENDPOINT = f"http://127.0.0.1:{APP_PORT}/chat/"
def create_new_session() -> int:
data = {"persona_id": 0} # Global default Persona
response = requests.post(LOCAL_CHAT_ENDPOINT + "create-chat-session", json=data)
response.raise_for_status()
new_session_id = response.json()["chat_session_id"]
return new_session_id
def send_chat_message(
message: str,
chat_session_id: int,
parent_message: int | None,
) -> int:
data = {
"message": message,
"chat_session_id": chat_session_id,
"parent_message_id": parent_message,
"prompt_id": 0, # Global default Prompt
"retrieval_options": {
"run_search": "always",
"real_time": True,
"filters": {"tags": []},
},
}
docs: list[dict] | None = None
message_id: int | None = None
with requests.post(
LOCAL_CHAT_ENDPOINT + "send-message", json=data, stream=True
) as r:
for json_response in r.iter_lines():
response_text = json.loads(json_response.decode())
new_token = response_text.get("answer_piece")
if docs is None:
docs = response_text.get("top_documents")
if message_id is None:
message_id = response_text.get("message_id")
if new_token:
print(new_token, end="", flush=True)
print()
if docs:
docs.sort(key=lambda x: x["score"], reverse=True) # type: ignore
print("\nReference Docs:")
for ind, doc in enumerate(docs, start=1):
print(f"\t - Doc {ind}: {doc.get('semantic_identifier')}")
if message_id is None:
raise ValueError("Couldn't get latest message id")
return message_id
def run_chat() -> None:
try:
new_session_id = create_new_session()
print(f"Chat Session ID: {new_session_id}")
except requests.exceptions.ConnectionError:
print(
"Looks like you haven't started the Danswer Backend server, please run the FastAPI server"
)
exit()
return
parent_message = None
while True:
new_message = input(
"\n\n----------------------------------\n"
"Please provide a new chat message:\n> "
)
parent_message = send_chat_message(
new_message, new_session_id, parent_message=parent_message
)
if __name__ == "__main__":
run_chat()