Chat Backend (#801)

backend/scripts/simulate_chat_frontend.py

@@ -1,11 +1,8 @@
 # This file is purely for development use, not included in any builds
-# Use this to test the chat feature with and without context.
-# With context refers to being able to call out to Danswer and other tools (currently no other tools)
-# Without context refers to only knowing the chat's own history with no additional information
+# 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 argparse
 import json
 
 import requests
@@ -16,7 +13,8 @@ LOCAL_CHAT_ENDPOINT = f"http://127.0.0.1:{APP_PORT}/chat/"
 
 
 def create_new_session() -> int:
-    response = requests.post(LOCAL_CHAT_ENDPOINT + "create-chat-session")
+    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
@@ -25,19 +23,18 @@ def create_new_session() -> int:
 def send_chat_message(
     message: str,
     chat_session_id: int,
-    message_number: int,
-    parent_edit_number: int | None,
-    persona_id: int | None,
-) -> None:
+    parent_message: int | None,
+) -> int:
     data = {
         "message": message,
         "chat_session_id": chat_session_id,
-        "message_number": message_number,
-        "parent_edit_number": parent_edit_number,
-        "persona_id": persona_id,
+        "parent_message_id": parent_message,
+        "prompt_id": 0,  # Global default Prompt
+        "retrieval_options": {"run_search": "always", "real_time": True},
     }
 
     docs: list[dict] | None = None
+    message_id: int | None = None
     with requests.post(
         LOCAL_CHAT_ENDPOINT + "send-message", json=data, stream=True
     ) as r:
@@ -46,17 +43,25 @@ def send_chat_message(
             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")
 
-def run_chat(contextual: bool) -> None:
+    return message_id
+
+
+def run_chat() -> None:
     try:
         new_session_id = create_new_session()
         print(f"Chat Session ID: {new_session_id}")
@@ -65,34 +70,19 @@ def run_chat(contextual: bool) -> None:
             "Looks like you haven't started the Danswer Backend server, please run the FastAPI server"
         )
         exit()
+        return
 
-    persona_id = 1 if contextual else None
-
-    message_num = 0
-    parent_edit = None
+    parent_message = None
     while True:
         new_message = input(
             "\n\n----------------------------------\n"
             "Please provide a new chat message:\n> "
         )
 
-        send_chat_message(
-            new_message, new_session_id, message_num, parent_edit, persona_id
+        parent_message = send_chat_message(
+            new_message, new_session_id, parent_message=parent_message
         )
 
-        message_num += 2  # 1 for User message, 1 for AI response
-        parent_edit = 0  # Since no edits, the parent message is always the first edit of that message number
-
 
 if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    parser.add_argument(
-        "-c",
-        "--contextual",
-        action="store_true",
-        help="If this flag is set, the chat is able to use retrieval",
-    )
-    args = parser.parse_args()
-
-    contextual = args.contextual
-    run_chat(contextual)
+    run_chat()