mirror of
https://github.com/DrHo1y/ezrknn-llm.git
synced 2026-10-02 08:06:37 +07:00
79 lines
3.5 KiB
Python
79 lines
3.5 KiB
Python
import sys
|
|
import requests
|
|
import json
|
|
|
|
# Set the address of the Server.
|
|
server_url = 'http://172.16.10.79:8080/rkllm_chat'
|
|
# Set whether to enable streaming mode.
|
|
is_streaming = True
|
|
|
|
# Create a session object.
|
|
session = requests.Session()
|
|
session.keep_alive = False # Close the connection pool to maintain a long connection.
|
|
adapter = requests.adapters.HTTPAdapter(max_retries=5)
|
|
session.mount('https://', adapter)
|
|
session.mount('http://', adapter)
|
|
|
|
if __name__ == '__main__':
|
|
print("============================")
|
|
print("Input your question in the terminal to start a conversation with the RKLLM model...")
|
|
print("============================")
|
|
# Enter a loop to continuously get user input and converse with the RKLLM model.
|
|
while True:
|
|
try:
|
|
user_message = input("\n*Please enter your question:")
|
|
if user_message == "exit":
|
|
print("============================")
|
|
print("The RKLLM Server is stopping......")
|
|
print("============================")
|
|
break
|
|
else:
|
|
# Set the request headers; in this case, the headers have no actual effect and are only used to simulate the OpenAI interface design.
|
|
headers = {
|
|
'Content-Type': 'application/json',
|
|
'Authorization': 'not_required'
|
|
}
|
|
|
|
# Prepare the data to be sent
|
|
# model: The model defined by the user when setting up RKLLM-Server; this has no effect here
|
|
# messages: The user's input question, which RKLLM-Server will use as input and return the model's reply; multiple questions can be added to messages
|
|
# stream: Whether to enable streaming conversation, similar to the OpenAI interface
|
|
data = {
|
|
"model": 'your_model_deploy_with_RKLLM_Server',
|
|
"messages": [{"role": "user", "content": user_message}],
|
|
"stream": is_streaming
|
|
}
|
|
|
|
# Send a POST request
|
|
responses = session.post(server_url, json=data, headers=headers, stream=is_streaming, verify=False)
|
|
|
|
if not is_streaming:
|
|
# Parse the response
|
|
if responses.status_code == 200:
|
|
print("Q:", data["messages"][-1]["content"])
|
|
print("A:", json.loads(responses.text)["choices"][-1]["message"]["content"])
|
|
else:
|
|
print("Error:", responses.text)
|
|
else:
|
|
if responses.status_code == 200:
|
|
print("Q:", data["messages"][-1]["content"])
|
|
print("A:", end="")
|
|
for line in responses.iter_lines():
|
|
if line:
|
|
line = json.loads(line.decode('utf-8'))
|
|
if line["choices"][-1]["finish_reason"] != "stop":
|
|
print(line["choices"][-1]["delta"]["content"], end="")
|
|
sys.stdout.flush()
|
|
else:
|
|
print('Error:', responses.text)
|
|
|
|
except KeyboardInterrupt:
|
|
# Capture Ctrl-C signal to close the session
|
|
session.close()
|
|
|
|
print("\n")
|
|
print("============================")
|
|
print("The RKLLM Server is stopping......")
|
|
print("============================")
|
|
break
|