521 lines
19 KiB
Python
521 lines
19 KiB
Python
#!/usr/bin/env python3
|
|
"""Streaming AI Bridge server for Jibo.
|
|
|
|
Replicates the original Jibo server pattern: responds sentence-by-sentence
|
|
so the robot can start speaking immediately while the LLM is still generating.
|
|
|
|
Endpoints:
|
|
- POST /v1/chat/text {"text": "..."} -> chunked NDJSON, one line per sentence:
|
|
{"sentence": "First sentence.", "done": false}
|
|
{"sentence": "Second sentence!", "done": false}
|
|
{"sentence": "", "done": true, "reply": "First sentence. Second sentence!"}
|
|
Legacy (non-streaming) clients still work: the final line has the full "reply".
|
|
|
|
- POST /v1/chat/audio {"wav_base64": "...", "sample_rate": 16000} -> same streaming format
|
|
|
|
LLM:
|
|
- Uses Ollama Chat API with streaming: http://localhost:11434/api/chat
|
|
Env:
|
|
OLLAMA_URL (default: http://127.0.0.1:11434/api/chat)
|
|
OLLAMA_MODEL (default: smollm2:135m)
|
|
|
|
STT (optional, for /audio):
|
|
- If `faster-whisper` is installed, it will be used.
|
|
Env:
|
|
WHISPER_MODEL (default: base)
|
|
|
|
Run:
|
|
python3 server.py --host 0.0.0.0 --port 8020
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import array
|
|
import base64
|
|
import io
|
|
import json
|
|
import os
|
|
import tempfile
|
|
import time
|
|
import traceback
|
|
import wave
|
|
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
|
from urllib.error import URLError
|
|
from urllib.request import Request, urlopen
|
|
|
|
|
|
def _ts() -> str:
|
|
# ISO-ish timestamp in local time; good enough for debugging
|
|
return time.strftime("%Y-%m-%d %H:%M:%S")
|
|
|
|
|
|
def _log(msg: str):
|
|
# Server previously muted logs; we want visibility while debugging.
|
|
print(f"[{_ts()}] {msg}", flush=True)
|
|
|
|
|
|
def _json_response(handler: BaseHTTPRequestHandler, status: int, payload: dict):
|
|
body = json.dumps(payload).encode("utf-8")
|
|
handler.send_response(status)
|
|
handler.send_header("Content-Type", "application/json")
|
|
handler.send_header("Content-Length", str(len(body)))
|
|
handler.end_headers()
|
|
handler.wfile.write(body)
|
|
|
|
|
|
def _read_json(handler: BaseHTTPRequestHandler) -> dict:
|
|
length = int(handler.headers.get("Content-Length", "0"))
|
|
raw = handler.rfile.read(length) if length else b"{}"
|
|
return json.loads(raw.decode("utf-8"))
|
|
|
|
|
|
def _ollama_chat(user_text: str) -> str:
|
|
"""Non-streaming fallback (used by /audio endpoint)."""
|
|
ollama_url = os.environ.get("OLLAMA_URL", "http://127.0.0.1:11434/api/chat")
|
|
model = os.environ.get("OLLAMA_MODEL", "smollm2:135m")
|
|
|
|
req_body = {
|
|
"model": model,
|
|
"stream": False,
|
|
"messages": [
|
|
{"role": "system", "content": "You are Jibo, a friendly home robot. Keep replies short and spoken."},
|
|
{"role": "user", "content": user_text},
|
|
],
|
|
}
|
|
|
|
req = Request(
|
|
ollama_url,
|
|
data=json.dumps(req_body).encode("utf-8"),
|
|
headers={"Content-Type": "application/json"},
|
|
method="POST",
|
|
)
|
|
|
|
try:
|
|
with urlopen(req, timeout=60) as resp:
|
|
data = json.loads(resp.read().decode("utf-8"))
|
|
except Exception as e:
|
|
_log(f"Ollama request failed url={ollama_url!r} err={e!r}")
|
|
raise
|
|
|
|
msg = data.get("message") or {}
|
|
content = msg.get("content")
|
|
if not content:
|
|
raise RuntimeError(f"Unexpected Ollama response: {data}")
|
|
return content.strip()
|
|
|
|
|
|
import re
|
|
|
|
# Sentence boundary: split on .!? followed by space or end, but not on
|
|
# abbreviations like "Dr." or "Mr." or decimals like "3.5".
|
|
_SENTENCE_END_RE = re.compile(
|
|
r'(?<=[.!?])\s+(?=[A-Z"\'])|(?<=[.!?])$'
|
|
)
|
|
|
|
|
|
def _split_sentences(text: str) -> list[str]:
|
|
"""Split text into sentences. Returns list of sentence strings."""
|
|
parts = _SENTENCE_END_RE.split(text.strip())
|
|
return [p.strip() for p in parts if p.strip()]
|
|
|
|
|
|
def _ollama_chat_streaming(user_text: str, on_sentence):
|
|
"""Stream Ollama response token-by-token, calling on_sentence(sentence_str)
|
|
each time a complete sentence is detected. Returns the full reply string.
|
|
|
|
This is the key to sub-5-second perceived latency: the robot starts speaking
|
|
the first sentence while the LLM is still generating the rest.
|
|
"""
|
|
ollama_url = os.environ.get("OLLAMA_URL", "http://127.0.0.1:11434/api/chat")
|
|
model = os.environ.get("OLLAMA_MODEL", "smollm2:135m")
|
|
|
|
req_body = {
|
|
"model": model,
|
|
"stream": True,
|
|
"messages": [
|
|
{"role": "system", "content": "You are Jibo, a friendly home robot. Keep replies short and spoken."},
|
|
{"role": "user", "content": user_text},
|
|
],
|
|
}
|
|
|
|
req = Request(
|
|
ollama_url,
|
|
data=json.dumps(req_body).encode("utf-8"),
|
|
headers={"Content-Type": "application/json"},
|
|
method="POST",
|
|
)
|
|
|
|
full_reply = ""
|
|
buffer = ""
|
|
sentences_sent = 0
|
|
t0 = time.monotonic()
|
|
|
|
with urlopen(req, timeout=120) as resp:
|
|
# Ollama streams one JSON object per line
|
|
for raw_line in resp:
|
|
line = raw_line.decode("utf-8", errors="replace").strip()
|
|
if not line:
|
|
continue
|
|
try:
|
|
chunk = json.loads(line)
|
|
except json.JSONDecodeError:
|
|
continue
|
|
|
|
# Each chunk: {"message": {"content": "token"}, "done": false}
|
|
msg = chunk.get("message") or {}
|
|
token = msg.get("content") or ""
|
|
if token:
|
|
buffer += token
|
|
full_reply += token
|
|
|
|
# Check if buffer contains a complete sentence to flush
|
|
# Look for sentence-ending punctuation followed by a space or more tokens
|
|
sentences = _split_sentences(buffer)
|
|
if len(sentences) > 1:
|
|
# All but the last are complete sentences; last is still building
|
|
for s in sentences[:-1]:
|
|
elapsed = time.monotonic() - t0
|
|
_log(f" sentence #{sentences_sent} at {elapsed:.2f}s: {s[:120]!r}")
|
|
on_sentence(s)
|
|
sentences_sent += 1
|
|
buffer = sentences[-1]
|
|
|
|
if chunk.get("done"):
|
|
break
|
|
|
|
# Flush any remaining text as the final sentence
|
|
leftover = buffer.strip()
|
|
if leftover:
|
|
elapsed = time.monotonic() - t0
|
|
_log(f" sentence #{sentences_sent} (final) at {elapsed:.2f}s: {leftover[:120]!r}")
|
|
on_sentence(leftover)
|
|
|
|
return full_reply.strip()
|
|
|
|
|
|
def _short_err(e: BaseException) -> str:
|
|
s = str(e) or e.__class__.__name__
|
|
s = " ".join(s.split())
|
|
if len(s) > 240:
|
|
s = s[:240] + "..."
|
|
return s
|
|
|
|
|
|
def _ollama_down_reply() -> str:
|
|
# Keep it short and speakable.
|
|
return "My AI server isn't reachable right now. Please start Ollama on the computer, then try again."
|
|
|
|
|
|
def _wav_diagnostics(wav_bytes: bytes) -> dict:
|
|
"""Best-effort WAV parsing + signal stats for debugging mic capture."""
|
|
info: dict = {"bytes": len(wav_bytes)}
|
|
try:
|
|
with wave.open(io.BytesIO(wav_bytes), "rb") as wf: # type: ignore[name-defined]
|
|
nch = wf.getnchannels()
|
|
sw = wf.getsampwidth()
|
|
fr = wf.getframerate()
|
|
nframes = wf.getnframes()
|
|
info.update({"channels": nch, "sample_width": sw, "frame_rate": fr, "frames": nframes})
|
|
|
|
# Read up to ~3 seconds of audio for stats (avoid huge CPU)
|
|
max_frames = min(nframes, fr * 3)
|
|
frames = wf.readframes(max_frames)
|
|
except Exception as e:
|
|
info["parse_error"] = str(e)
|
|
return info
|
|
|
|
# Only compute stats for 16-bit PCM (most common).
|
|
if info.get("sample_width") != 2:
|
|
return info
|
|
|
|
try:
|
|
samples = array.array("h")
|
|
samples.frombytes(frames)
|
|
if not samples:
|
|
return info
|
|
|
|
mn = min(samples)
|
|
mx = max(samples)
|
|
zeros = sum(1 for s in samples if s == 0)
|
|
# RMS over interleaved samples (good enough for quick signal presence)
|
|
n = float(len(samples))
|
|
rms = (sum(float(s) * float(s) for s in samples) / n) ** 0.5
|
|
# Per-channel RMS (helps debug mic arrays)
|
|
ch = int(info.get("channels") or 1)
|
|
channel_rms = None
|
|
if ch > 1:
|
|
channel_rms = []
|
|
for c in range(ch):
|
|
chan = samples[c::ch]
|
|
if not chan:
|
|
channel_rms.append(0.0)
|
|
else:
|
|
nn = float(len(chan))
|
|
channel_rms.append((sum(float(s) * float(s) for s in chan) / nn) ** 0.5)
|
|
info.update(
|
|
{
|
|
"min": int(mn),
|
|
"max": int(mx),
|
|
"rms": float(rms),
|
|
"zero_frac": float(zeros) / n,
|
|
"channel_rms": channel_rms,
|
|
}
|
|
)
|
|
except Exception as e:
|
|
info["stats_error"] = str(e)
|
|
return info
|
|
|
|
|
|
def _to_loudest_channel_mono_wav(wav_bytes: bytes) -> tuple[bytes, dict]:
|
|
"""If WAV is multi-channel 16-bit PCM, pick loudest channel and return mono WAV bytes."""
|
|
try:
|
|
with wave.open(io.BytesIO(wav_bytes), "rb") as wf:
|
|
nch = wf.getnchannels()
|
|
sw = wf.getsampwidth()
|
|
fr = wf.getframerate()
|
|
nframes = wf.getnframes()
|
|
frames = wf.readframes(nframes)
|
|
except Exception as e:
|
|
return wav_bytes, {"convert_error": str(e)}
|
|
|
|
if nch <= 1 or sw != 2:
|
|
return wav_bytes, {"converted": False, "channels": nch, "sample_width": sw}
|
|
|
|
samples = array.array("h")
|
|
samples.frombytes(frames)
|
|
if not samples:
|
|
return wav_bytes, {"converted": False, "reason": "empty_samples"}
|
|
|
|
# Choose loudest channel by RMS
|
|
rms_list: list[float] = []
|
|
for c in range(nch):
|
|
chan = samples[c::nch]
|
|
if not chan:
|
|
rms_list.append(0.0)
|
|
else:
|
|
nn = float(len(chan))
|
|
rms_list.append((sum(float(s) * float(s) for s in chan) / nn) ** 0.5)
|
|
best = max(range(nch), key=lambda i: rms_list[i])
|
|
mono = samples[best::nch]
|
|
|
|
out = io.BytesIO()
|
|
with wave.open(out, "wb") as ow:
|
|
ow.setnchannels(1)
|
|
ow.setsampwidth(2)
|
|
ow.setframerate(fr)
|
|
ow.writeframes(mono.tobytes())
|
|
|
|
return out.getvalue(), {"converted": True, "picked_channel": best, "channel_rms": rms_list, "frame_rate": fr}
|
|
|
|
|
|
class _Whisper:
|
|
def __init__(self):
|
|
self._model = None
|
|
|
|
def available(self) -> bool:
|
|
try:
|
|
import faster_whisper # noqa: F401
|
|
|
|
return True
|
|
except Exception:
|
|
return False
|
|
|
|
def transcribe_wav_bytes(self, wav_bytes: bytes) -> str:
|
|
try:
|
|
from faster_whisper import WhisperModel
|
|
except Exception as e:
|
|
raise RuntimeError(
|
|
"Audio mode requires `faster-whisper` (pip install faster-whisper) and ffmpeg on your PC"
|
|
) from e
|
|
|
|
model_name = os.environ.get("WHISPER_MODEL", "base")
|
|
if self._model is None:
|
|
# CPU-friendly default; user can override via WHISPER_MODEL and faster-whisper params if needed.
|
|
self._model = WhisperModel(model_name, device="cpu", compute_type="int8")
|
|
|
|
with tempfile.NamedTemporaryFile(suffix=".wav", delete=True) as f:
|
|
f.write(wav_bytes)
|
|
f.flush()
|
|
segments, info = self._model.transcribe(f.name)
|
|
text = "".join(seg.text for seg in segments).strip()
|
|
return text
|
|
|
|
|
|
_whisper = _Whisper()
|
|
|
|
|
|
class Handler(BaseHTTPRequestHandler):
|
|
server_version = "JiboAIBridge/1.0"
|
|
|
|
def do_POST(self):
|
|
try:
|
|
client = f"{self.client_address[0]}:{self.client_address[1]}"
|
|
length = int(self.headers.get("Content-Length", "0") or "0")
|
|
_log(f"{client} POST {self.path} len={length}")
|
|
|
|
if self.path == "/v1/chat/text":
|
|
payload = _read_json(self)
|
|
text = (payload.get("text") or "").strip()
|
|
if not text:
|
|
_json_response(self, 400, {"error": "Missing 'text'"})
|
|
return
|
|
_log(f"{client} /text prompt_chars={len(text)} prompt={text[:200]!r}")
|
|
try:
|
|
# Stream response: chunked transfer encoding with NDJSON
|
|
# Each line is a JSON object the robot can parse immediately.
|
|
# This replicates the original Jibo hub pattern where the robot
|
|
# starts acting on partial results while the server is still working.
|
|
self.send_response(200)
|
|
self.send_header("Content-Type", "application/x-ndjson")
|
|
self.send_header("Transfer-Encoding", "chunked")
|
|
self.end_headers()
|
|
|
|
t0 = time.monotonic()
|
|
|
|
def _send_chunk(obj: dict):
|
|
line = json.dumps(obj) + "\n"
|
|
data = line.encode("utf-8")
|
|
# HTTP chunked encoding: hex size + CRLF + data + CRLF
|
|
self.wfile.write(f"{len(data):x}\r\n".encode())
|
|
self.wfile.write(data)
|
|
self.wfile.write(b"\r\n")
|
|
self.wfile.flush()
|
|
|
|
def _on_sentence(sentence: str):
|
|
_send_chunk({"sentence": sentence, "done": False})
|
|
|
|
reply = _ollama_chat_streaming(text, _on_sentence)
|
|
elapsed = time.monotonic() - t0
|
|
_log(f"{client} /text ok reply_chars={len(reply)} elapsed={elapsed:.2f}s")
|
|
|
|
# Final chunk with full reply (for logging / legacy compat)
|
|
_send_chunk({"sentence": "", "done": True, "reply": reply})
|
|
# Terminate chunked encoding
|
|
self.wfile.write(b"0\r\n\r\n")
|
|
self.wfile.flush()
|
|
|
|
except (URLError, ConnectionRefusedError) as e:
|
|
_log(f"{client} /text ollama_unreachable err={_short_err(e)!r}")
|
|
# Fallback: try to send a non-streaming error response.
|
|
# If headers already sent, write it as a chunk.
|
|
try:
|
|
err_reply = _ollama_down_reply()
|
|
err_obj = {"sentence": err_reply, "done": True, "reply": err_reply,
|
|
"ollama_ok": False, "ollama_error": _short_err(e)}
|
|
line = json.dumps(err_obj) + "\n"
|
|
data = line.encode("utf-8")
|
|
self.wfile.write(f"{len(data):x}\r\n".encode())
|
|
self.wfile.write(data)
|
|
self.wfile.write(b"\r\n0\r\n\r\n")
|
|
self.wfile.flush()
|
|
except Exception:
|
|
pass # headers may not have been sent yet
|
|
return
|
|
|
|
if self.path == "/v1/chat/audio":
|
|
payload = _read_json(self)
|
|
b64 = payload.get("wav_base64")
|
|
if not b64:
|
|
_json_response(self, 400, {"error": "Missing 'wav_base64'"})
|
|
return
|
|
|
|
if not _whisper.available():
|
|
_log(f"{client} /audio STT unavailable (faster-whisper not installed)")
|
|
_json_response(
|
|
self,
|
|
503,
|
|
{
|
|
"error": "STT unavailable: install faster-whisper and ffmpeg on this PC",
|
|
"hint": "pip install faster-whisper (and install ffmpeg)",
|
|
},
|
|
)
|
|
return
|
|
|
|
wav_bytes = base64.b64decode(b64)
|
|
diag = _wav_diagnostics(wav_bytes)
|
|
wav_for_stt, conv = _to_loudest_channel_mono_wav(wav_bytes)
|
|
if conv.get("converted"):
|
|
diag_mono = _wav_diagnostics(wav_for_stt)
|
|
_log(
|
|
f"{client} /audio wav_diag={json.dumps(diag, sort_keys=True)} "
|
|
f"mono={json.dumps(diag_mono, sort_keys=True)} conv={json.dumps(conv, sort_keys=True)}"
|
|
)
|
|
else:
|
|
_log(f"{client} /audio wav_diag={json.dumps(diag, sort_keys=True)} conv={json.dumps(conv, sort_keys=True)}")
|
|
try:
|
|
# Save for debugging (overwrite each time)
|
|
out_path = os.environ.get("AI_BRIDGE_LAST_WAV", "jibo_last.wav")
|
|
with open(out_path, "wb") as f:
|
|
f.write(wav_bytes)
|
|
_log(f"{client} /audio decoded bytes={len(wav_bytes)} saved={out_path}")
|
|
except Exception as e:
|
|
_log(f"{client} /audio failed saving wav: {e}")
|
|
|
|
transcript = _whisper.transcribe_wav_bytes(wav_for_stt)
|
|
if not transcript:
|
|
_log(f"{client} /audio empty transcript")
|
|
_json_response(self, 200, {"reply": "I didn't catch that. Could you say it again?", "text": ""})
|
|
return
|
|
_log(f"{client} /audio transcript_chars={len(transcript)} transcript={transcript[:200]!r}")
|
|
try:
|
|
reply = _ollama_chat(transcript)
|
|
_log(f"{client} /audio ok text_chars={len(transcript)} reply_chars={len(reply)}")
|
|
_json_response(self, 200, {"reply": reply, "text": transcript})
|
|
except URLError as e:
|
|
_log(f"{client} /audio ollama_unreachable err={_short_err(e)!r}")
|
|
_json_response(
|
|
self,
|
|
200,
|
|
{"reply": _ollama_down_reply(), "text": transcript, "ollama_ok": False, "ollama_error": _short_err(e)},
|
|
)
|
|
except ConnectionRefusedError as e:
|
|
_log(f"{client} /audio ollama_refused err={_short_err(e)!r}")
|
|
_json_response(
|
|
self,
|
|
200,
|
|
{"reply": _ollama_down_reply(), "text": transcript, "ollama_ok": False, "ollama_error": _short_err(e)},
|
|
)
|
|
return
|
|
|
|
_json_response(self, 404, {"error": "Not found"})
|
|
|
|
except Exception as e:
|
|
_log(f"ERROR {self.path}: {e}\n{traceback.format_exc()}")
|
|
_json_response(
|
|
self,
|
|
500,
|
|
{
|
|
"error": str(e),
|
|
"trace": traceback.format_exc(),
|
|
},
|
|
)
|
|
|
|
def log_message(self, format, *args):
|
|
# Keep BaseHTTPRequestHandler from double-logging; we do our own.
|
|
return
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--host", default="0.0.0.0")
|
|
ap.add_argument("--port", type=int, default=8020)
|
|
args = ap.parse_args()
|
|
|
|
server = ThreadingHTTPServer((args.host, args.port), Handler)
|
|
_log(f"AI Bridge server listening on http://{args.host}:{args.port}")
|
|
_log(
|
|
"Ollama: "
|
|
+ os.environ.get("OLLAMA_URL", "http://127.0.0.1:11434/api/chat")
|
|
+ " model="
|
|
+ os.environ.get("OLLAMA_MODEL", "smollm2:135m")
|
|
)
|
|
_log("Ollama health check: curl -s http://127.0.0.1:11434/api/tags | head")
|
|
if not _whisper.available():
|
|
_log("STT: faster-whisper not installed; /v1/chat/audio will return 503")
|
|
server.serve_forever()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|