forked from onyx-dot-app/onyx
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathapi_inference_sample.py
86 lines (69 loc) · 2.88 KB
/
api_inference_sample.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
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
)