Initial commit
This commit is contained in:
nikolya
@@ -0,0 +1,827 @@
|
||||
import base64
|
||||
import json
|
||||
import os
|
||||
import sqlite3
|
||||
import tempfile
|
||||
import threading
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
from typing import Any, Dict, Generator, Optional, Tuple
|
||||
|
||||
import requests
|
||||
from flask import Flask, Response, jsonify, request, send_from_directory
|
||||
|
||||
|
||||
APP_DIR = os.path.dirname(os.path.abspath(__file__))
|
||||
CONFIG_PATH = os.path.join(APP_DIR, "config.json")
|
||||
DB_PATH = os.path.join(APP_DIR, "chat.sqlite3")
|
||||
_CONFIG_LOCK = threading.Lock()
|
||||
|
||||
DEFAULT_ENDPOINTS = {
|
||||
"ollama": "http://localhost:11434",
|
||||
# "ollama-web" in this app means "Ollama Cloud via ollama.com API"
|
||||
"ollama-web": "https://ollama.com",
|
||||
}
|
||||
|
||||
|
||||
def _b64e(s: str) -> str:
|
||||
return base64.b64encode(s.encode("utf-8")).decode("utf-8")
|
||||
|
||||
|
||||
def _b64d(s: str) -> str:
|
||||
return base64.b64decode(s.encode("utf-8")).decode("utf-8")
|
||||
|
||||
|
||||
def _mask_key(k: str) -> str:
|
||||
if not k:
|
||||
return ""
|
||||
if len(k) <= 6:
|
||||
return "*" * len(k)
|
||||
return f"{k[:3]}***{k[-3:]}"
|
||||
|
||||
|
||||
def load_config() -> Dict[str, Any]:
|
||||
if not os.path.exists(CONFIG_PATH):
|
||||
return {"providers": {}}
|
||||
try:
|
||||
with open(CONFIG_PATH, "r", encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
if not isinstance(data, dict):
|
||||
return {"providers": {}}
|
||||
data.setdefault("providers", {})
|
||||
if not isinstance(data["providers"], dict):
|
||||
data["providers"] = {}
|
||||
return data
|
||||
except Exception:
|
||||
return {"providers": {}}
|
||||
|
||||
|
||||
def save_config(cfg: Dict[str, Any]) -> None:
|
||||
# Use a unique temp file to avoid races between concurrent requests.
|
||||
# Keep temp file in the same directory to make os.replace atomic.
|
||||
with _CONFIG_LOCK:
|
||||
fd, tmp = tempfile.mkstemp(prefix="config.", suffix=".json.tmp", dir=APP_DIR)
|
||||
try:
|
||||
with os.fdopen(fd, "w", encoding="utf-8") as f:
|
||||
json.dump(cfg, f, ensure_ascii=False, indent=2)
|
||||
f.flush()
|
||||
os.fsync(f.fileno())
|
||||
os.replace(tmp, CONFIG_PATH)
|
||||
finally:
|
||||
try:
|
||||
if os.path.exists(tmp):
|
||||
os.unlink(tmp)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def get_provider_cfg(provider: str) -> Dict[str, Any]:
|
||||
cfg = load_config()
|
||||
prov = cfg.get("providers", {}).get(provider, {})
|
||||
return prov if isinstance(prov, dict) else {}
|
||||
|
||||
|
||||
def set_provider_cfg(provider: str, provider_cfg: Dict[str, Any]) -> None:
|
||||
cfg = load_config()
|
||||
cfg.setdefault("providers", {})
|
||||
cfg["providers"][provider] = provider_cfg
|
||||
save_config(cfg)
|
||||
|
||||
|
||||
def get_api_key(provider: str) -> str:
|
||||
prov = get_provider_cfg(provider)
|
||||
enc = prov.get("api_key_b64", "")
|
||||
if not enc:
|
||||
return ""
|
||||
try:
|
||||
return _b64d(enc)
|
||||
except Exception:
|
||||
return ""
|
||||
|
||||
|
||||
def get_endpoint(provider: str) -> str:
|
||||
prov = get_provider_cfg(provider)
|
||||
ep = prov.get("endpoint", "")
|
||||
if isinstance(ep, str) and ep.strip():
|
||||
return ep.strip().rstrip("/")
|
||||
return DEFAULT_ENDPOINTS.get(provider, "")
|
||||
|
||||
|
||||
def json_line(obj: Dict[str, Any]) -> bytes:
|
||||
return (json.dumps(obj, ensure_ascii=False) + "\n").encode("utf-8")
|
||||
|
||||
|
||||
def _decode_line(raw: Any) -> str:
|
||||
if raw is None:
|
||||
return ""
|
||||
if isinstance(raw, bytes):
|
||||
return raw.decode("utf-8", errors="replace")
|
||||
if isinstance(raw, str):
|
||||
return raw
|
||||
return str(raw)
|
||||
|
||||
|
||||
def _extract_ollama_library_models_from_text(txt: str) -> list:
|
||||
# Web page fetch (no stable public JSON API). Parse headings like "## llama3.2".
|
||||
models = []
|
||||
for line in (txt or "").splitlines():
|
||||
s = line.strip()
|
||||
if not s.startswith("## "):
|
||||
continue
|
||||
name = s[3:].strip()
|
||||
if name:
|
||||
models.append(name)
|
||||
uniq = sorted(set(models))
|
||||
return uniq
|
||||
|
||||
|
||||
def _extract_ollama_library_tags_from_model_page(txt: str, base_model: str) -> list:
|
||||
"""
|
||||
Parse a model library page (converted to text/markdown by WebFetch-like rendering)
|
||||
and extract available tags like "<base_model>:<tag>".
|
||||
"""
|
||||
tags: set[str] = set()
|
||||
needle = f"/library/{base_model}:"
|
||||
for line in (txt or "").splitlines():
|
||||
if needle not in line:
|
||||
continue
|
||||
# Example lines in fetched text:
|
||||
# [llama3.2:latest2.0GB · ...](https://ollama.com/library/llama3.2:latest)
|
||||
# [llama3.2:1b1.3GB · ...](https://ollama.com/library/llama3.2:1b)
|
||||
i = line.find(needle)
|
||||
if i < 0:
|
||||
continue
|
||||
frag = line[i + len("/library/") :]
|
||||
# frag now starts with "<base_model>:<tag>..."
|
||||
if not frag.startswith(base_model + ":"):
|
||||
continue
|
||||
frag2 = frag[len(base_model) + 1 :]
|
||||
tag = ""
|
||||
for ch in frag2:
|
||||
if ch.isalnum() or ch in ("-", "_", "."):
|
||||
tag += ch
|
||||
else:
|
||||
break
|
||||
if tag:
|
||||
tags.add(f"{base_model}:{tag}")
|
||||
if not tags:
|
||||
# Fallback: at least a usable default
|
||||
return [f"{base_model}:latest"]
|
||||
return sorted(tags)
|
||||
|
||||
|
||||
def _ollama_cloud_headers() -> Dict[str, str]:
|
||||
key = get_api_key("ollama-web")
|
||||
if not key:
|
||||
raise RuntimeError("Ollama.com API key is missing (create one at https://ollama.com/settings/keys)")
|
||||
return {"Authorization": f"Bearer {key}"}
|
||||
|
||||
|
||||
def _stream_ollama_ndjson(r: requests.Response) -> Generator[Tuple[str, str], None, None]:
|
||||
"""
|
||||
Ollama streaming endpoints can return application/x-ndjson where each line is a JSON object.
|
||||
We emit "content" deltas based on response/message fields.
|
||||
"""
|
||||
for raw in r.iter_lines(decode_unicode=False):
|
||||
line = _decode_line(raw).strip()
|
||||
if not line:
|
||||
continue
|
||||
try:
|
||||
obj = json.loads(line)
|
||||
except Exception:
|
||||
continue
|
||||
|
||||
if not isinstance(obj, dict):
|
||||
continue
|
||||
|
||||
# /api/generate shape: {"response": "...", "done": false, ...}
|
||||
resp = obj.get("response")
|
||||
if isinstance(resp, str) and resp:
|
||||
yield "content", resp
|
||||
continue
|
||||
|
||||
# /api/chat shape: {"message": {"role": "...", "content": "..."}, "done": false, ...}
|
||||
msg = obj.get("message")
|
||||
if isinstance(msg, dict):
|
||||
content = msg.get("content")
|
||||
if isinstance(content, str) and content:
|
||||
yield "content", content
|
||||
continue
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ProviderResult:
|
||||
ok: bool
|
||||
error: Optional[str] = None
|
||||
|
||||
|
||||
def _required_str(payload: Dict[str, Any], key: str) -> Tuple[Optional[str], Optional[str]]:
|
||||
v = payload.get(key)
|
||||
if not isinstance(v, str) or not v.strip():
|
||||
return None, f"Missing or invalid field: {key}"
|
||||
return v, None
|
||||
|
||||
|
||||
app = Flask(__name__, static_folder=".", static_url_path="")
|
||||
|
||||
|
||||
def _db_connect() -> sqlite3.Connection:
|
||||
conn = sqlite3.connect(DB_PATH, timeout=30, check_same_thread=False)
|
||||
conn.row_factory = sqlite3.Row
|
||||
try:
|
||||
conn.execute("PRAGMA journal_mode=WAL;")
|
||||
except Exception:
|
||||
# Some environments can reject WAL; keep default mode.
|
||||
pass
|
||||
conn.execute("PRAGMA foreign_keys=ON;")
|
||||
return conn
|
||||
|
||||
|
||||
def init_db() -> None:
|
||||
conn = _db_connect()
|
||||
try:
|
||||
conn.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS chat_messages (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
provider TEXT NOT NULL,
|
||||
model TEXT NOT NULL,
|
||||
prompt TEXT NOT NULL,
|
||||
response TEXT NOT NULL DEFAULT '',
|
||||
reasoning TEXT NOT NULL DEFAULT '',
|
||||
reasoning_enabled INTEGER NOT NULL DEFAULT 0,
|
||||
status TEXT NOT NULL DEFAULT 'streaming',
|
||||
error TEXT NOT NULL DEFAULT '',
|
||||
created_at INTEGER NOT NULL,
|
||||
updated_at INTEGER NOT NULL
|
||||
);
|
||||
"""
|
||||
)
|
||||
conn.execute("CREATE INDEX IF NOT EXISTS idx_chat_provider_id ON chat_messages(provider, id);")
|
||||
conn.commit()
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
def _now_ts() -> int:
|
||||
return int(time.time())
|
||||
|
||||
|
||||
def _chat_create(provider: str, model: str, prompt: str, reasoning_enabled: bool) -> int:
|
||||
ts = _now_ts()
|
||||
conn = _db_connect()
|
||||
try:
|
||||
cur = conn.execute(
|
||||
"""
|
||||
INSERT INTO chat_messages(provider, model, prompt, reasoning_enabled, status, created_at, updated_at)
|
||||
VALUES(?, ?, ?, ?, 'streaming', ?, ?)
|
||||
""",
|
||||
(provider, model, prompt, 1 if reasoning_enabled else 0, ts, ts),
|
||||
)
|
||||
conn.commit()
|
||||
return int(cur.lastrowid)
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
def _chat_update(
|
||||
chat_id: int,
|
||||
*,
|
||||
response: Optional[str] = None,
|
||||
reasoning: Optional[str] = None,
|
||||
status: Optional[str] = None,
|
||||
error: Optional[str] = None,
|
||||
) -> None:
|
||||
ts = _now_ts()
|
||||
sets = ["updated_at = ?"]
|
||||
args: list[Any] = [ts]
|
||||
if response is not None:
|
||||
sets.append("response = ?")
|
||||
args.append(response)
|
||||
if reasoning is not None:
|
||||
sets.append("reasoning = ?")
|
||||
args.append(reasoning)
|
||||
if status is not None:
|
||||
sets.append("status = ?")
|
||||
args.append(status)
|
||||
if error is not None:
|
||||
sets.append("error = ?")
|
||||
args.append(error)
|
||||
args.append(chat_id)
|
||||
|
||||
conn = _db_connect()
|
||||
try:
|
||||
conn.execute(f"UPDATE chat_messages SET {', '.join(sets)} WHERE id = ?", args)
|
||||
conn.commit()
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
# Ensure DB exists even when run via `flask run` / WSGI.
|
||||
init_db()
|
||||
|
||||
|
||||
@app.get("/")
|
||||
def index() -> Any:
|
||||
return send_from_directory(APP_DIR, "index.html")
|
||||
|
||||
|
||||
@app.get("/api/config")
|
||||
def api_get_config() -> Any:
|
||||
cfg = load_config()
|
||||
out: Dict[str, Any] = {"providers": {}}
|
||||
for provider, p in cfg.get("providers", {}).items():
|
||||
if not isinstance(p, dict):
|
||||
continue
|
||||
key = ""
|
||||
try:
|
||||
key = _b64d(p.get("api_key_b64", "") or "")
|
||||
except Exception:
|
||||
key = ""
|
||||
out["providers"][provider] = {
|
||||
"hasKey": bool(key),
|
||||
"keyMasked": _mask_key(key) if key else "",
|
||||
"endpoint": p.get("endpoint", DEFAULT_ENDPOINTS.get(provider, "")),
|
||||
"lastModel": p.get("last_model", ""),
|
||||
}
|
||||
return jsonify(out)
|
||||
|
||||
|
||||
@app.post("/api/config")
|
||||
def api_set_config() -> Any:
|
||||
payload = request.get_json(silent=True) or {}
|
||||
provider, err = _required_str(payload, "provider")
|
||||
if err:
|
||||
return jsonify({"ok": False, "error": err}), 400
|
||||
|
||||
api_key = payload.get("apiKey", "")
|
||||
if api_key is not None and not isinstance(api_key, str):
|
||||
return jsonify({"ok": False, "error": "apiKey must be a string"}), 400
|
||||
|
||||
endpoint = payload.get("endpoint", None)
|
||||
if endpoint is not None and not isinstance(endpoint, str):
|
||||
return jsonify({"ok": False, "error": "endpoint must be a string"}), 400
|
||||
|
||||
last_model = payload.get("lastModel", None)
|
||||
if last_model is not None and not isinstance(last_model, str):
|
||||
return jsonify({"ok": False, "error": "lastModel must be a string"}), 400
|
||||
|
||||
current = get_provider_cfg(provider)
|
||||
new_cfg = dict(current)
|
||||
|
||||
if isinstance(api_key, str) and api_key.strip():
|
||||
new_cfg["api_key_b64"] = _b64e(api_key.strip())
|
||||
elif payload.get("clearKey") is True:
|
||||
new_cfg.pop("api_key_b64", None)
|
||||
|
||||
if isinstance(endpoint, str):
|
||||
new_cfg["endpoint"] = endpoint.strip().rstrip("/")
|
||||
|
||||
if isinstance(last_model, str) and last_model.strip():
|
||||
new_cfg["last_model"] = last_model.strip()
|
||||
|
||||
set_provider_cfg(provider, new_cfg)
|
||||
return jsonify({"ok": True})
|
||||
|
||||
|
||||
@app.get("/api/models")
|
||||
def api_models() -> Any:
|
||||
provider = request.args.get("provider", "").strip()
|
||||
if not provider:
|
||||
return jsonify({"ok": False, "error": "Missing provider"}), 400
|
||||
|
||||
try:
|
||||
models = discover_models(provider)
|
||||
return jsonify({"ok": True, "models": models})
|
||||
except requests.exceptions.ConnectionError:
|
||||
if provider == "ollama":
|
||||
return (
|
||||
jsonify(
|
||||
{
|
||||
"ok": False,
|
||||
"error": "Ollama не запущена. Проверьте, что локальный Ollama работает на http://localhost:11434.",
|
||||
}
|
||||
),
|
||||
502,
|
||||
)
|
||||
return jsonify({"ok": False, "error": "Network error while loading models"}), 502
|
||||
except Exception as e:
|
||||
return jsonify({"ok": False, "error": str(e) or "Failed to load models"}), 500
|
||||
|
||||
|
||||
def discover_models(provider: str) -> list:
|
||||
if provider == "ollama":
|
||||
base = get_endpoint("ollama") or DEFAULT_ENDPOINTS["ollama"]
|
||||
r = requests.get(f"{base}/api/tags", timeout=10)
|
||||
r.raise_for_status()
|
||||
data = r.json()
|
||||
models = []
|
||||
for m in data.get("models", []) or []:
|
||||
name = m.get("name")
|
||||
if isinstance(name, str) and name:
|
||||
models.append(name)
|
||||
return sorted(set(models))
|
||||
|
||||
if provider == "ollama-web":
|
||||
# Ollama Cloud: https://ollama.com/api/tags (requires API key)
|
||||
base = get_endpoint("ollama-web") or DEFAULT_ENDPOINTS["ollama-web"]
|
||||
r = requests.get(f"{base}/api/tags", headers=_ollama_cloud_headers(), timeout=20)
|
||||
r.raise_for_status()
|
||||
data = r.json()
|
||||
models = []
|
||||
for m in data.get("models", []) or []:
|
||||
if not isinstance(m, dict):
|
||||
continue
|
||||
name = m.get("name")
|
||||
if isinstance(name, str) and name:
|
||||
models.append(name)
|
||||
return sorted(set(models))
|
||||
|
||||
if provider == "openrouter":
|
||||
key = get_api_key("openrouter")
|
||||
if not key:
|
||||
raise RuntimeError("OpenRouter API key is missing")
|
||||
r = requests.get(
|
||||
"https://openrouter.ai/api/v1/models",
|
||||
headers={"Authorization": f"Bearer {key}"},
|
||||
timeout=20,
|
||||
)
|
||||
r.raise_for_status()
|
||||
data = r.json()
|
||||
models = []
|
||||
for m in data.get("data", []) or []:
|
||||
mid = m.get("id")
|
||||
if isinstance(mid, str) and mid:
|
||||
models.append(mid)
|
||||
return sorted(set(models))
|
||||
|
||||
if provider == "nvidia":
|
||||
key = get_api_key("nvidia")
|
||||
if not key:
|
||||
raise RuntimeError("NVIDIA API key is missing")
|
||||
r = requests.get(
|
||||
"https://integrate.api.nvidia.com/v1/models",
|
||||
headers={"Authorization": f"Bearer {key}", "Accept": "application/json"},
|
||||
timeout=20,
|
||||
)
|
||||
r.raise_for_status()
|
||||
data = r.json()
|
||||
models = []
|
||||
items = []
|
||||
if isinstance(data, dict):
|
||||
if isinstance(data.get("data"), list):
|
||||
items = data["data"]
|
||||
elif isinstance(data.get("models"), list):
|
||||
items = data["models"]
|
||||
for m in items or []:
|
||||
if not isinstance(m, dict):
|
||||
continue
|
||||
mid = m.get("id") or m.get("name")
|
||||
if isinstance(mid, str) and mid:
|
||||
models.append(mid)
|
||||
return sorted(set(models))
|
||||
|
||||
if provider == "huggingface":
|
||||
# List only models that are actually served by at least one Inference Provider.
|
||||
# This avoids many 404s on https://api-inference.huggingface.co/models/{id}.
|
||||
key = get_api_key("huggingface")
|
||||
headers = {"Authorization": f"Bearer {key}"} if key else {}
|
||||
r = requests.get(
|
||||
"https://huggingface.co/api/models"
|
||||
"?inference_provider=all"
|
||||
"&pipeline_tag=text-generation"
|
||||
"&sort=downloads"
|
||||
"&direction=-1"
|
||||
"&limit=80",
|
||||
headers=headers,
|
||||
timeout=20,
|
||||
)
|
||||
r.raise_for_status()
|
||||
data = r.json()
|
||||
models = []
|
||||
if isinstance(data, list):
|
||||
for m in data:
|
||||
if isinstance(m, dict) and isinstance(m.get("modelId"), str):
|
||||
models.append(m["modelId"])
|
||||
return sorted(set(models))
|
||||
|
||||
raise RuntimeError(f"Unknown provider: {provider}")
|
||||
|
||||
|
||||
@app.get("/api/chat")
|
||||
def api_chat_list() -> Any:
|
||||
provider = request.args.get("provider", "").strip()
|
||||
if not provider:
|
||||
return jsonify({"ok": False, "error": "Missing provider"}), 400
|
||||
|
||||
try:
|
||||
limit = int(request.args.get("limit", "20"))
|
||||
except Exception:
|
||||
limit = 20
|
||||
limit = max(1, min(limit, 100))
|
||||
|
||||
before_id_raw = request.args.get("before_id", "").strip()
|
||||
before_id: Optional[int] = None
|
||||
if before_id_raw:
|
||||
try:
|
||||
before_id = int(before_id_raw)
|
||||
except Exception:
|
||||
return jsonify({"ok": False, "error": "before_id must be an integer"}), 400
|
||||
|
||||
conn = _db_connect()
|
||||
try:
|
||||
if before_id is None:
|
||||
rows = conn.execute(
|
||||
"""
|
||||
SELECT id, provider, model, prompt, response, reasoning, reasoning_enabled, status, error, created_at, updated_at
|
||||
FROM chat_messages
|
||||
WHERE provider = ?
|
||||
ORDER BY id DESC
|
||||
LIMIT ?
|
||||
""",
|
||||
(provider, limit),
|
||||
).fetchall()
|
||||
else:
|
||||
rows = conn.execute(
|
||||
"""
|
||||
SELECT id, provider, model, prompt, response, reasoning, reasoning_enabled, status, error, created_at, updated_at
|
||||
FROM chat_messages
|
||||
WHERE provider = ? AND id < ?
|
||||
ORDER BY id DESC
|
||||
LIMIT ?
|
||||
""",
|
||||
(provider, before_id, limit),
|
||||
).fetchall()
|
||||
|
||||
items = []
|
||||
for r in rows:
|
||||
items.append(
|
||||
{
|
||||
"id": int(r["id"]),
|
||||
"provider": r["provider"],
|
||||
"model": r["model"],
|
||||
"prompt": r["prompt"],
|
||||
"response": r["response"],
|
||||
"reasoning": r["reasoning"],
|
||||
"reasoningEnabled": bool(r["reasoning_enabled"]),
|
||||
"status": r["status"],
|
||||
"error": r["error"],
|
||||
"createdAt": int(r["created_at"]),
|
||||
"updatedAt": int(r["updated_at"]),
|
||||
}
|
||||
)
|
||||
|
||||
next_before_id = items[-1]["id"] if items else None
|
||||
has_more = len(items) == limit
|
||||
return jsonify({"ok": True, "items": items, "nextBeforeId": next_before_id, "hasMore": has_more})
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
@app.post("/api/chat/clear")
|
||||
def api_chat_clear() -> Any:
|
||||
payload = request.get_json(silent=True) or {}
|
||||
provider, err = _required_str(payload, "provider")
|
||||
if err:
|
||||
return jsonify({"ok": False, "error": err}), 400
|
||||
conn = _db_connect()
|
||||
try:
|
||||
conn.execute("DELETE FROM chat_messages WHERE provider = ?", (provider,))
|
||||
conn.commit()
|
||||
return jsonify({"ok": True})
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
@app.post("/api/query")
|
||||
def api_query() -> Any:
|
||||
payload = request.get_json(silent=True) or {}
|
||||
provider, err = _required_str(payload, "provider")
|
||||
if err:
|
||||
return jsonify({"ok": False, "error": err}), 400
|
||||
model, err = _required_str(payload, "model")
|
||||
if err:
|
||||
return jsonify({"ok": False, "error": err}), 400
|
||||
prompt, err = _required_str(payload, "prompt")
|
||||
if err:
|
||||
return jsonify({"ok": False, "error": err}), 400
|
||||
reasoning = bool(payload.get("reasoning", False))
|
||||
|
||||
# persist last model
|
||||
try:
|
||||
set_provider_cfg(provider, {**get_provider_cfg(provider), "last_model": model})
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
chat_id: Optional[int] = None
|
||||
try:
|
||||
chat_id = _chat_create(provider, model, prompt, reasoning_enabled=reasoning)
|
||||
except Exception:
|
||||
chat_id = None
|
||||
|
||||
def gen() -> Generator[bytes, None, None]:
|
||||
yield json_line(
|
||||
{"type": "meta", "provider": provider, "model": model, "ts": int(time.time()), "chatId": chat_id}
|
||||
)
|
||||
full_text = ""
|
||||
full_reasoning = ""
|
||||
try:
|
||||
for channel, chunk in stream_completion(provider, model, prompt, reasoning=reasoning):
|
||||
if channel == "reasoning":
|
||||
full_reasoning += chunk
|
||||
else:
|
||||
full_text += chunk
|
||||
yield json_line({"type": "delta", "channel": channel, "text": chunk})
|
||||
if chat_id is not None:
|
||||
_chat_update(chat_id, response=full_text, reasoning=full_reasoning, status="done", error="")
|
||||
yield json_line({"type": "done"})
|
||||
except requests.exceptions.ConnectionError:
|
||||
if chat_id is not None:
|
||||
_chat_update(
|
||||
chat_id,
|
||||
response=full_text,
|
||||
reasoning=full_reasoning,
|
||||
status="error",
|
||||
error="Network error",
|
||||
)
|
||||
yield json_line({"type": "error", "error": "Network error. Проверьте соединение и повторите."})
|
||||
except Exception as e:
|
||||
msg = str(e) or "Request failed"
|
||||
if chat_id is not None:
|
||||
_chat_update(chat_id, response=full_text, reasoning=full_reasoning, status="error", error=msg)
|
||||
yield json_line({"type": "error", "error": msg})
|
||||
except GeneratorExit:
|
||||
if chat_id is not None:
|
||||
try:
|
||||
_chat_update(chat_id, response=full_text, reasoning=full_reasoning, status="aborted")
|
||||
except Exception:
|
||||
pass
|
||||
raise
|
||||
|
||||
return Response(gen(), mimetype="application/x-ndjson")
|
||||
|
||||
|
||||
def stream_completion(
|
||||
provider: str, model: str, prompt: str, reasoning: bool
|
||||
) -> Generator[Tuple[str, str], None, None]:
|
||||
if provider == "ollama":
|
||||
base = get_endpoint("ollama") or DEFAULT_ENDPOINTS["ollama"]
|
||||
r = requests.post(
|
||||
f"{base}/api/generate",
|
||||
json={"model": model, "prompt": prompt, "stream": True},
|
||||
timeout=60,
|
||||
stream=True,
|
||||
)
|
||||
r.raise_for_status()
|
||||
for raw in r.iter_lines(decode_unicode=False):
|
||||
line = _decode_line(raw).strip()
|
||||
if not line:
|
||||
continue
|
||||
try:
|
||||
obj = json.loads(line)
|
||||
except Exception:
|
||||
continue
|
||||
if isinstance(obj, dict) and isinstance(obj.get("response"), str):
|
||||
yield "content", obj["response"]
|
||||
return
|
||||
|
||||
if provider == "ollama-web":
|
||||
# Ollama Cloud via ollama.com API: https://ollama.com/api/chat
|
||||
base = get_endpoint("ollama-web") or DEFAULT_ENDPOINTS["ollama-web"]
|
||||
headers = {"Content-Type": "application/json", **_ollama_cloud_headers()}
|
||||
body: Dict[str, Any] = {
|
||||
"model": model,
|
||||
"stream": True,
|
||||
"messages": [{"role": "user", "content": prompt}],
|
||||
}
|
||||
if reasoning:
|
||||
body["reasoning"] = True
|
||||
r = requests.post(f"{base}/api/chat", headers=headers, json=body, timeout=60, stream=True)
|
||||
r.raise_for_status()
|
||||
yield from _stream_ollama_ndjson(r)
|
||||
return
|
||||
|
||||
if provider == "openrouter":
|
||||
key = get_api_key("openrouter")
|
||||
if not key:
|
||||
raise RuntimeError("OpenRouter API key is missing")
|
||||
headers = {
|
||||
"Authorization": f"Bearer {key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
body: Dict[str, Any] = {
|
||||
"model": model,
|
||||
"stream": True,
|
||||
"messages": [{"role": "user", "content": prompt}],
|
||||
}
|
||||
if reasoning:
|
||||
body["reasoning"] = True
|
||||
r = requests.post(
|
||||
"https://openrouter.ai/api/v1/chat/completions",
|
||||
headers=headers,
|
||||
json=body,
|
||||
timeout=60,
|
||||
stream=True,
|
||||
)
|
||||
r.raise_for_status()
|
||||
yield from _stream_openai_sse(r)
|
||||
return
|
||||
|
||||
if provider == "nvidia":
|
||||
key = get_api_key("nvidia")
|
||||
if not key:
|
||||
raise RuntimeError("NVIDIA API key is missing")
|
||||
headers = {
|
||||
"Authorization": f"Bearer {key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
body: Dict[str, Any] = {
|
||||
"model": model,
|
||||
"stream": True,
|
||||
"messages": [{"role": "user", "content": prompt}],
|
||||
}
|
||||
if reasoning:
|
||||
body["reasoning"] = True
|
||||
r = requests.post(
|
||||
"https://integrate.api.nvidia.com/v1/chat/completions",
|
||||
headers=headers,
|
||||
json=body,
|
||||
timeout=60,
|
||||
stream=True,
|
||||
)
|
||||
r.raise_for_status()
|
||||
yield from _stream_openai_sse(r)
|
||||
return
|
||||
|
||||
if provider == "huggingface":
|
||||
key = get_api_key("huggingface")
|
||||
if not key:
|
||||
raise RuntimeError("HuggingFace API key is missing")
|
||||
headers = {"Authorization": f"Bearer {key}"}
|
||||
# Basic inference API. Not truly streaming; emulate.
|
||||
r = requests.post(
|
||||
f"https://api-inference.huggingface.co/models/{model}",
|
||||
headers=headers,
|
||||
json={"inputs": prompt, "parameters": {"max_new_tokens": 512}},
|
||||
timeout=90,
|
||||
)
|
||||
if r.status_code == 404:
|
||||
raise RuntimeError(
|
||||
"HuggingFace: модель не найдена или недоступна для Inference API (404). "
|
||||
"Попробуйте другую модель из списка (text-generation) или проверьте права доступа/гейтинг."
|
||||
)
|
||||
r.raise_for_status()
|
||||
data = r.json()
|
||||
text = ""
|
||||
if isinstance(data, list) and data and isinstance(data[0], dict) and isinstance(data[0].get("generated_text"), str):
|
||||
text = data[0]["generated_text"]
|
||||
elif isinstance(data, dict) and isinstance(data.get("generated_text"), str):
|
||||
text = data["generated_text"]
|
||||
elif isinstance(data, dict) and isinstance(data.get("error"), str):
|
||||
raise RuntimeError(data["error"])
|
||||
else:
|
||||
text = json.dumps(data, ensure_ascii=False)
|
||||
for ch in text:
|
||||
yield "content", ch
|
||||
return
|
||||
|
||||
raise RuntimeError(f"Unknown provider: {provider}")
|
||||
|
||||
|
||||
def _stream_openai_sse(r: requests.Response) -> Generator[Tuple[str, str], None, None]:
|
||||
for raw in r.iter_lines(decode_unicode=False):
|
||||
line = _decode_line(raw).strip()
|
||||
if not line:
|
||||
continue
|
||||
if not line.startswith("data:"):
|
||||
continue
|
||||
data = line[len("data:") :].strip()
|
||||
if data == "[DONE]":
|
||||
break
|
||||
try:
|
||||
obj = json.loads(data)
|
||||
except Exception:
|
||||
continue
|
||||
try:
|
||||
delta = obj["choices"][0].get("delta", {})
|
||||
content = delta.get("content")
|
||||
if isinstance(content, str) and content:
|
||||
yield "content", content
|
||||
|
||||
reasoning = delta.get("reasoning") or delta.get("thinking")
|
||||
if isinstance(reasoning, str) and reasoning:
|
||||
yield "reasoning", reasoning
|
||||
except Exception:
|
||||
continue
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# In some sandboxed environments, Werkzeug's debugger can fail due to /dev/shm restrictions.
|
||||
debug = os.environ.get("FLASK_DEBUG", "").strip() in ("1", "true", "True")
|
||||
app.run(
|
||||
host="0.0.0.0",
|
||||
port=int(os.environ.get("PORT", "8000")),
|
||||
debug=debug,
|
||||
use_reloader=debug,
|
||||
)
|
||||
Reference in New Issue
Block a user