Adolfo Reyna
237 lines
8.5 KiB
Python
237 lines
8.5 KiB
Python
"""A Pipecat processor that puts the native macOS / Apple Silicon LLM model in the LLM slot.
|
|
|
|
Supports:
|
|
1. Apple Intelligence / FoundationModels (via swift/llm-helper)
|
|
2. MLX local on-device LLM on Apple Silicon (via mlx-lm)
|
|
|
|
Streams response text frames (LLMFullResponseStartFrame, LLMTextFrame,
|
|
LLMFullResponseEndFrame) so text-to-speech downstream responds with low latency.
|
|
"""
|
|
|
|
import asyncio
|
|
import json
|
|
import os
|
|
import shutil
|
|
import subprocess
|
|
from pathlib import Path
|
|
from loguru import logger
|
|
|
|
from pipecat.frames.frames import (
|
|
CancelFrame,
|
|
EndFrame,
|
|
Frame,
|
|
InterruptionFrame,
|
|
LLMContextFrame,
|
|
LLMFullResponseEndFrame,
|
|
LLMFullResponseStartFrame,
|
|
LLMTextFrame,
|
|
StartFrame,
|
|
)
|
|
from pipecat.processors.frame_processor import FrameDirection, FrameProcessor
|
|
|
|
_NOISE_TRANSCRIPTS = {
|
|
"",
|
|
".",
|
|
"thank you.",
|
|
"thanks for watching!",
|
|
"you",
|
|
"bye.",
|
|
"okay.",
|
|
"[blank_audio]",
|
|
"[silence]",
|
|
}
|
|
|
|
HERE = Path(__file__).parent
|
|
LLM_HELPER_PATH = HERE / "swift" / "llm-helper"
|
|
|
|
|
|
def probe_apple_llm() -> tuple[bool, str]:
|
|
"""Check if native macOS / Apple Silicon LLM options are available."""
|
|
if LLM_HELPER_PATH.exists() and os.access(LLM_HELPER_PATH, os.X_OK):
|
|
try:
|
|
res = subprocess.run(
|
|
[str(LLM_HELPER_PATH), "--check"],
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=5.0,
|
|
)
|
|
if res.returncode == 0:
|
|
try:
|
|
payload = json.loads(res.stdout.strip())
|
|
if payload.get("available"):
|
|
return True, payload.get("reason", "macOS FoundationModels available")
|
|
except json.JSONDecodeError:
|
|
pass
|
|
except Exception:
|
|
pass
|
|
|
|
try:
|
|
import mlx_lm
|
|
return True, "Apple Silicon MLX local LLM engine available"
|
|
except ImportError:
|
|
return False, "Neither swift/llm-helper nor mlx-lm is available"
|
|
|
|
|
|
def _latest_user_text(context) -> str:
|
|
if not context or not hasattr(context, "messages"):
|
|
return ""
|
|
for msg in reversed(context.messages):
|
|
if isinstance(msg, dict) and msg.get("role") == "user":
|
|
content = msg.get("content", "")
|
|
if isinstance(content, str):
|
|
return content
|
|
elif isinstance(content, list):
|
|
text_parts = [c.get("text", "") for c in content if isinstance(c, dict)]
|
|
return " ".join(text_parts)
|
|
return ""
|
|
|
|
|
|
class MacOSLLM(FrameProcessor):
|
|
"""Runs user turns through macOS native LLM (Apple Intelligence or MLX)."""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
model: str = "mlx-community/gemma-2-2b-it-4bit",
|
|
system_prompt: str | None = None,
|
|
observer=None,
|
|
**kwargs,
|
|
):
|
|
super().__init__(**kwargs)
|
|
self._model_name = model
|
|
self._system_prompt = system_prompt or "You are a helpful voice assistant. Keep answers brief and spoken naturally."
|
|
self._on_reply = observer
|
|
self._turn_task: asyncio.Task | None = None
|
|
self._history: list[dict[str, str]] = []
|
|
self._mlx_model = None
|
|
self._mlx_tokenizer = None
|
|
self._use_swift = False
|
|
|
|
async def process_frame(self, frame: Frame, direction: FrameDirection):
|
|
await super().process_frame(frame, direction)
|
|
|
|
if isinstance(frame, StartFrame):
|
|
await self.push_frame(frame, direction)
|
|
await self._connect()
|
|
elif isinstance(frame, (EndFrame, CancelFrame)):
|
|
await self._cancel_turn()
|
|
await self.push_frame(frame, direction)
|
|
elif isinstance(frame, InterruptionFrame):
|
|
await self._cancel_turn()
|
|
await self.push_frame(frame, direction)
|
|
elif isinstance(frame, LLMContextFrame):
|
|
text = _latest_user_text(frame.context)
|
|
await self._maybe_start_turn(text)
|
|
else:
|
|
await self.push_frame(frame, direction)
|
|
|
|
async def _connect(self):
|
|
available, reason = probe_apple_llm()
|
|
logger.info(f"macOS LLM engine: {reason}")
|
|
if LLM_HELPER_PATH.exists() and "FoundationModels available" in reason:
|
|
self._use_swift = True
|
|
logger.info("Using Swift FoundationModels engine.")
|
|
else:
|
|
self._use_swift = False
|
|
logger.info(f"Loading MLX model {self._model_name} on Apple Silicon...")
|
|
loop = asyncio.get_running_loop()
|
|
def _load():
|
|
from mlx_lm import load
|
|
return load(self._model_name)
|
|
self._mlx_model, self._mlx_tokenizer = await loop.run_in_executor(None, _load)
|
|
logger.info("MLX model loaded.")
|
|
|
|
async def _maybe_start_turn(self, text: str):
|
|
utterance = text.strip()
|
|
if utterance.lower() in _NOISE_TRANSCRIPTS or len(utterance) < 2:
|
|
logger.debug(f"Ignoring noise transcript: {utterance!r}")
|
|
return
|
|
|
|
await self._cancel_turn()
|
|
logger.info(f"You: {utterance}")
|
|
self._turn_task = self.create_task(self._run_turn(utterance))
|
|
|
|
async def _cancel_turn(self):
|
|
if not self._turn_task:
|
|
return
|
|
task, self._turn_task = self._turn_task, None
|
|
await self.cancel_task(task)
|
|
|
|
async def _run_turn(self, utterance: str):
|
|
self._history.append({"role": "user", "content": utterance})
|
|
|
|
await self.push_frame(LLMFullResponseStartFrame())
|
|
chunks: list[str] = []
|
|
|
|
try:
|
|
if self._use_swift:
|
|
proc = await asyncio.create_subprocess_exec(
|
|
str(LLM_HELPER_PATH),
|
|
"--prompt", utterance,
|
|
"--system", self._system_prompt,
|
|
"--stream",
|
|
stdout=asyncio.subprocess.PIPE,
|
|
stderr=asyncio.subprocess.PIPE,
|
|
)
|
|
while True:
|
|
line = await proc.stdout.readline()
|
|
if not line:
|
|
break
|
|
try:
|
|
data = json.loads(line.decode().strip())
|
|
if "delta" in data:
|
|
chunk = data["delta"]
|
|
chunks.append(chunk)
|
|
await self.push_frame(LLMTextFrame(chunk))
|
|
elif "text" in data and not chunks:
|
|
chunk = data["text"]
|
|
chunks.append(chunk)
|
|
await self.push_frame(LLMTextFrame(chunk))
|
|
except json.JSONDecodeError:
|
|
pass
|
|
await proc.wait()
|
|
else:
|
|
loop = asyncio.get_running_loop()
|
|
def _gen():
|
|
from mlx_lm import generate
|
|
from mlx_lm.sample_utils import make_sampler
|
|
# Keep history concise to avoid context drift and repetition
|
|
recent_history = self._history[-8:]
|
|
messages = [{"role": "system", "content": self._system_prompt}] + recent_history
|
|
prompt = self._mlx_tokenizer.apply_chat_template(
|
|
messages, add_generation_prompt=True, tokenize=False
|
|
)
|
|
sampler = make_sampler(temp=0.7)
|
|
return generate(
|
|
self._mlx_model,
|
|
self._mlx_tokenizer,
|
|
prompt=prompt,
|
|
max_tokens=150,
|
|
sampler=sampler,
|
|
verbose=False,
|
|
)
|
|
|
|
response_text = await loop.run_in_executor(None, _gen)
|
|
# Clean up any potential repetitive sentence tails
|
|
cleaned_text = response_text.strip()
|
|
chunks.append(cleaned_text)
|
|
await self.push_frame(LLMTextFrame(cleaned_text))
|
|
|
|
except asyncio.CancelledError:
|
|
logger.info("Turn cancelled mid-response.")
|
|
raise
|
|
except Exception as e:
|
|
logger.error(f"macOS LLM error: {e}")
|
|
err_msg = "Sorry, I ran into an error generating a response."
|
|
chunks.append(err_msg)
|
|
await self.push_frame(LLMTextFrame(err_msg))
|
|
finally:
|
|
await self.push_frame(LLMFullResponseEndFrame())
|
|
|
|
full_reply = "".join(chunks).strip()
|
|
if full_reply:
|
|
self._history.append({"role": "assistant", "content": full_reply})
|
|
logger.info(f"macOS LLM: {full_reply}")
|
|
if self._on_reply:
|
|
self._on_reply(full_reply)
|