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 functools
import logging
import signal
import threading

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


def build_triage(cfg, client, on_result=None):
    """Build (analyzer, lane, brief_fn, brief_reset); (None, None, None, None) when disabled."""
    if not cfg.triage.enabled:
        return (None, 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,
    )
    brief_fn = None
    brief_reset = None
    if client is not None:
        engine = ConversationEngine(
            cfg.llm,
            client,
            system_prompt=cfg.llm.triage_system_prompt,
            compaction_system_prompt=cfg.llm.brief_compaction_system_prompt,
            max_tokens=cfg.triage.extract_max_tokens,
            timeout_s=cfg.triage.extract_timeout_s,
            slot_id=cfg.triage.brief_slot_id,
            notify=None,
        )
        brief_fn = functools.partial(engine.send, finalize=analyzer.finalize_brief)
        brief_reset = engine.reset

    return (analyzer, lane, brief_fn, brief_reset)


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 = ConversationEngine(
            cfg.llm,
            client,
            system_prompt=cfg.llm.system_prompt_chat,
            slot_id=cfg.llm.chat_slot_id,
        )
        executor = LlmExecutor(
            cfg.llm,
            client,
            default_registry(),
            slot_id=cfg.llm.command_slot_id,
        )
        chat_fn = session.send
        chat_reset = session.reset

    holder = {}
    triage, lane, brief_fn, brief_reset = build_triage(
        cfg,
        client if cfg.llm.enabled else None,
        on_result=lambda gen, payload: holder["svc"]._on_triage_result(gen, payload),
    )
    brief_file = BriefFile(cfg.brief.write, cfg.brief.dir, cfg.brief.filename)
    _debug_factory = None
    debug_cfg = getattr(cfg, "debug", None)
    if debug_cfg is not None and bool(getattr(debug_cfg, "enabled", False)):
        try:
            from ..debug import DebugSink
        except Exception:  # noqa: BLE001 - debug stays disabled without the module
            DebugSink = None  # type: ignore[assignment]
        if DebugSink is not None:
            debug_dir = str(getattr(debug_cfg, "dir", "debug"))

            def _make_debug_sink():
                return DebugSink(debug_dir, "assistant")

            _debug_factory = _make_debug_sink

    service = AssistantService(
        cfg,
        subscriber,
        audit=audit,
        executor=executor,
        chat_fn=chat_fn,
        chat_reset=chat_reset,
        triage=triage,
        laya_lane=lane,
        brief_fn=brief_fn,
        brief_reset=brief_reset,
        brief_file=brief_file,
        debug_sink_factory=_debug_factory,
    )
    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