STT-Assistant/src/stt_daemon/assistant/runner.py

"""Run the assistant process: subscribe to the STT bus and serve the control socket."""

import argparse
import logging
import signal
import threading

from ..bus import EventSubscriber
from ..config import load_config
from .audit import AuditLog
from .control import ControlServer
from .executor import LlmExecutor, NullExecutor
from .llm import ChatSession, LlmClient
from .service import AssistantService
from .tools import default_registry


def build_triage(cfg, client, on_result=None):
    """Build (analyzer, lane, extract_fn); (None, None, None) when disabled."""
    if not cfg.triage.enabled:
        return (None, None, None)
    from .laya_lane import LayaLane
    from .triage import TriageAnalyzer

    analyzer = TriageAnalyzer(cfg.triage)
    lane = LayaLane(
        cfg.triage,
        on_state=lambda state, reason: analyzer.refresh(state, reason),
        on_result=on_result,
    )
    extract = None
    if client is not None:
        def _extract(prompt):
            resp = client.chat(
                [
                    {"role": "system", "content": cfg.llm.triage_system_prompt},
                    {"role": "user", "content": prompt},
                ],
                max_tokens=cfg.triage.extract_max_tokens,
                timeout_s=cfg.triage.extract_timeout_s,
            )
            return resp.content

        extract = _extract

    return (analyzer, lane, extract)


def main() -> int:
    """Entry point for the `stt-assistant` process."""
    parser = argparse.ArgumentParser(
        description="Assistant controller: subscribe to STT events and serve the control socket"
    )
    parser.add_argument("--config", type=str, default=None, help="Path to config.toml")
    parser.add_argument("-v", "--verbose", action="store_true", help="Enable debug logging")
    args = parser.parse_args()

    logging.basicConfig(
        level=logging.DEBUG if args.verbose else logging.INFO,
        format="[%(levelname)s] %(message)s",
    )

    cfg = load_config(args.config)
    subscriber = EventSubscriber(cfg.events.socket or None)
    audit = AuditLog(enabled=cfg.audit.enabled, directory=cfg.audit.dir)

    executor = NullExecutor()
    chat_fn = None
    chat_reset = None
    session = None
    client = None
    if cfg.llm.enabled:
        client = LlmClient(cfg.llm)
        session = ChatSession(cfg.llm, client)
        executor = LlmExecutor(cfg.llm, client, default_registry())
        chat_fn = session.send
        chat_reset = session.reset

    holder = {}
    triage, lane, extract_fn = build_triage(
        cfg,
        client if cfg.llm.enabled else None,
        on_result=lambda gen, payload: holder["svc"]._on_triage_result(gen, payload),
    )
    service = AssistantService(
        cfg,
        subscriber,
        audit=audit,
        executor=executor,
        chat_fn=chat_fn,
        chat_reset=chat_reset,
        triage=triage,
        laya_lane=lane,
        triage_extract_fn=extract_fn,
    )
    holder["svc"] = service
    if session is not None:
        session.notify = service.chat_notify
    control = ControlServer(cfg, service)
    service.set_sinks(emit=control.broadcast, relay=control.relay)

    control.start()
    service.start()
    if lane is not None:
        lane.start()
    logging.info("Assistant ready. Control socket: %s", control.path)

    stop = threading.Event()

    def handle_signal(_signum, _frame):
        stop.set()

    signal.signal(signal.SIGINT, handle_signal)
    signal.signal(signal.SIGTERM, handle_signal)
    stop.wait()

    service.stop()
    if lane is not None:
        lane.stop()
    control.close()
    return 0