107 lines
2.9 KiB
Python
107 lines
2.9 KiB
Python
import json
|
|
from dataclasses import asdict, dataclass
|
|
from urllib import request, error
|
|
|
|
from aiohttp import web
|
|
|
|
from server import PromptServer
|
|
|
|
OLLAMA_HOST = "http://127.0.0.1:11434"
|
|
OLLAMA_URL_GENERATE = f"{OLLAMA_HOST}/api/generate"
|
|
OLLAMA_URL_CHAT = f"{OLLAMA_HOST}/api/chat"
|
|
OLLAMA_TAGS_URL = f"{OLLAMA_HOST}/api/tags"
|
|
|
|
|
|
# Server api definitions
|
|
|
|
@PromptServer.instance.routes.get("/prompt_bench/ollama/models")
|
|
async def get_ollama_models(_request):
|
|
return web.json_response(_fetch_ollama_models())
|
|
|
|
|
|
# Custom request definitions
|
|
|
|
def _fetch_ollama_models() -> list[str]:
|
|
print("fetching ollama models")
|
|
try:
|
|
with request.urlopen(OLLAMA_TAGS_URL, timeout=5) as response:
|
|
print("received ollama models")
|
|
body = json.loads(response.read().decode("utf-8"))
|
|
models = [
|
|
model.get("name", "").strip()
|
|
for model in body.get("models", [])
|
|
if model.get("name", "").strip()
|
|
]
|
|
return models
|
|
except Exception:
|
|
print("fetching ollama models failed")
|
|
return []
|
|
|
|
|
|
def post_ollama_text(model: str, prompt: str):
|
|
payload = {
|
|
"model": model,
|
|
"prompt": prompt,
|
|
"stream": False,
|
|
}
|
|
|
|
req = request.Request(
|
|
OLLAMA_URL_GENERATE,
|
|
data=json.dumps(payload).encode("utf-8"),
|
|
headers={"Content-Type": "application/json"},
|
|
method="POST",
|
|
)
|
|
|
|
try:
|
|
with request.urlopen(req, timeout=120) as response:
|
|
body = json.loads(response.read().decode("utf-8"))
|
|
except error.URLError as exc:
|
|
print(exc)
|
|
# Ignore error
|
|
|
|
if "error" in body:
|
|
raise RuntimeError(f"Ollama generation failed: {body['error']}")
|
|
|
|
text = body.get("response", "").strip()
|
|
if not text:
|
|
raise RuntimeError("Ollama returned an empty response.")
|
|
|
|
return text
|
|
|
|
|
|
@dataclass
|
|
class OllamaChatMessage:
|
|
role: str
|
|
content: str
|
|
|
|
|
|
def post_ollama_chat(model: str, messages: list[OllamaChatMessage]):
|
|
payload = {
|
|
"model": model,
|
|
"messages": [asdict(m) for m in messages],
|
|
"stream": False,
|
|
}
|
|
|
|
req = request.Request(
|
|
OLLAMA_URL_CHAT,
|
|
data=json.dumps(payload).encode("utf-8"),
|
|
headers={"Content-Type": "application/json"},
|
|
method="POST",
|
|
)
|
|
|
|
try:
|
|
with request.urlopen(req, timeout=120) as response:
|
|
body = json.loads(response.read().decode("utf-8"))
|
|
except error.URLError as exc:
|
|
print(exc)
|
|
# Ignore error
|
|
|
|
if "error" in body:
|
|
raise RuntimeError(f"Ollama generation failed: {body['error']}")
|
|
|
|
assistant_message_raw = body.get("message", {"role": "assistant", "content": None})
|
|
if not assistant_message_raw.get("content", None):
|
|
raise RuntimeError("Ollama returned an empty response.")
|
|
|
|
return OllamaChatMessage(role="assistant", content=assistant_message_raw["content"])
|