90 lines
2.4 KiB
Python
90 lines
2.4 KiB
Python
|
import json
|
||
|
|
from functools import wraps
|
||
|
|
import requests
|
||
|
|
import sys
|
||
|
|
from qtpy.QtWidgets import QApplication, QDialog, QLineEdit, QPushButton, QVBoxLayout, QComboBox, QTextEdit
|
||
|
|
from qtpy.QtCore import QRunnable, QThreadPool
|
||
|
|
|
||
|
|
API_HOST = "http://localhost:11434/api"
|
||
|
|
|
||
|
|
|
||
|
|
def run_off_main_thread(func):
|
||
|
|
@wraps(func)
|
||
|
|
def wrapper(*args, **kwargs):
|
||
|
|
runner = QRunnable.create(lambda: func(*args, **kwargs))
|
||
|
|
QThreadPool.globalInstance().start(runner)
|
||
|
|
return wrapper
|
||
|
|
|
||
|
|
|
||
|
|
class Form(QDialog):
|
||
|
|
def __init__(self, parent=None):
|
||
|
|
super(Form, self).__init__(parent)
|
||
|
|
self.setWindowTitle("Ollama API")
|
||
|
|
|
||
|
|
self.layout = QVBoxLayout()
|
||
|
|
self.prompt = QLineEdit()
|
||
|
|
self.prompt.setPlaceholderText("Enter a prompt")
|
||
|
|
self.submit = QPushButton("Submit")
|
||
|
|
self.stream = QPushButton("Stream")
|
||
|
|
self.clear = QPushButton("Clear")
|
||
|
|
self.model_select = QComboBox()
|
||
|
|
self.response = QTextEdit()
|
||
|
|
|
||
|
|
self.layout.addWidget(self.prompt)
|
||
|
|
self.layout.addWidget(self.model_select)
|
||
|
|
self.layout.addWidget(self.submit)
|
||
|
|
self.layout.addWidget(self.stream)
|
||
|
|
self.layout.addWidget(self.clear)
|
||
|
|
self.layout.addWidget(self.response)
|
||
|
|
self.setLayout(self.layout)
|
||
|
|
|
||
|
|
self.get_models()
|
||
|
|
self.submit.clicked.connect(lambda: self.send_prompt())
|
||
|
|
self.stream.clicked.connect(lambda: self.stream_response())
|
||
|
|
self.clear.clicked.connect(lambda: self.clear_response())
|
||
|
|
|
||
|
|
# https://github.com/ollama/ollama/blob/main/docs/api.md#list-local-models
|
||
|
|
def get_models(self):
|
||
|
|
response = requests.get(f"{API_HOST}/tags")
|
||
|
|
print(response.json())
|
||
|
|
for models in response.json().get("models"):
|
||
|
|
self.model_select.addItem(models["model"])
|
||
|
|
|
||
|
|
def clear_response(self):
|
||
|
|
self.response.clear()
|
||
|
|
|
||
|
|
@run_off_main_thread
|
||
|
|
def send_prompt(self):
|
||
|
|
prompt = self.prompt.text()
|
||
|
|
model = self.model_select.currentText()
|
||
|
|
response = requests.post(f"{API_HOST}/generate", json={
|
||
|
|
"model": model,
|
||
|
|
"prompt": prompt,
|
||
|
|
"stream": False
|
||
|
|
})
|
||
|
|
print(response.json())
|
||
|
|
self.response.insertPlainText(response.json().get("response"))
|
||
|
|
|
||
|
|
@run_off_main_thread
|
||
|
|
def stream_response(self):
|
||
|
|
prompt = self.prompt.text()
|
||
|
|
model = self.model_select.currentText()
|
||
|
|
response = requests.post(f"{API_HOST}/generate", json={
|
||
|
|
"model": model,
|
||
|
|
"prompt": prompt,
|
||
|
|
"stream": True
|
||
|
|
}, stream=True)
|
||
|
|
for line in response.iter_lines():
|
||
|
|
if line:
|
||
|
|
text = json.loads(line)["response"]
|
||
|
|
print(text, end="")
|
||
|
|
self.response.insertPlainText(text)
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
app = QApplication(sys.argv)
|
||
|
|
form = Form()
|
||
|
|
form.show()
|
||
|
|
sys.exit(app.exec())
|