Sample API Script (onyx-dot-app#1079)

sernddev · Feb 13, 2024 · 13b2dd0 · 13b2dd0
1 parent 5925ecb
commit 13b2dd0
Showing 1 changed file with 86 additions and 0 deletions.
diff --git a/backend/scripts/api_inference_sample.py b/backend/scripts/api_inference_sample.py
@@ -0,0 +1,86 @@
+# This file is used to demonstrate how to use the backend APIs directly
+# In this case, the equivalent of asking a question in Danswer Chat in a new chat session
+import argparse
+import json
+import os
+
+import requests
+
+
+def create_new_chat_session(danswer_url: str, api_key: str | None) -> int:
+    headers = {"Authorization": f"Bearer {api_key}"} if api_key else None
+    session_endpoint = danswer_url + "/api/chat/create-chat-session"
+
+    response = requests.post(
+        session_endpoint,
+        headers=headers,
+        json={"persona_id": 0},  # Global default Persona/Assistant ID
+    )
+    response.raise_for_status()
+
+    new_session_id = response.json()["chat_session_id"]
+    return new_session_id
+
+
+def process_question(danswer_url: str, question: str, api_key: str | None) -> None:
+    message_endpoint = danswer_url + "/api/chat/send-message"
+
+    chat_session_id = create_new_chat_session(danswer_url, api_key)
+
+    headers = {"Authorization": f"Bearer {api_key}"} if api_key else None
+
+    data = {
+        "message": question,
+        "chat_session_id": chat_session_id,
+        "parent_message_id": None,
+        # Default Question Answer prompt
+        "prompt_id": 0,
+        # Not specifying any specific docs to chat to, we want to run a search
+        "search_doc_ids": None,
+        "retrieval_options": {
+            "run_search": "always",
+            "real_time": True,
+            "enable_auto_detect_filters": False,
+            # No filters applied, check all sources, document-sets, time ranges, etc.
+            "filters": {},
+        },
+    }
+
+    with requests.post(message_endpoint, headers=headers, json=data) as response:
+        response.raise_for_status()
+
+        for packet in response.iter_lines():
+            response_text = json.loads(packet.decode())
+            # Can also check "top_documents" to capture the streamed search results
+            # that include the highest matching documents to the query
+            # or check "message_id" to get the message_id used as parent_message_id
+            # to create follow-up messages
+            new_token = response_text.get("answer_piece")
+
+            if new_token:
+                print(new_token, end="", flush=True)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Sample API Usage")
+    parser.add_argument(
+        "--danswer-url",
+        type=str,
+        default="http://localhost:80",
+        help="Danswer URL, should point to Danswer nginx.",
+    )
+    parser.add_argument(
+        "--test-question",
+        type=str,
+        default="What is Danswer?",
+        help="Test question for new Chat Session.",
+    )
+
+    # Not needed if Auth is disabled
+    # Or for Danswer MIT API key must be replaced with session cookie
+    api_key = os.environ.get("DANSWER_API_KEY")
+
+    args = parser.parse_args()
+    process_question(
+        danswer_url=args.danswer_url, question=args.test_question, api_key=api_key
+    )