Modular async pipeline: FrameSource->Vision->context and STT/text->Brain->TTS. All stages are Protocols; mock backends run end-to-end with no deps/keys. Real backends included: mss screen capture, Claude vision+brain (guarded imports).
115 lines
3.6 KiB
Python
115 lines
3.6 KiB
Python
"""Orchestrator: wires the perception loop and the conversation loop together."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
|
|
from .interfaces import (
|
|
Brain,
|
|
FrameSource,
|
|
Reply,
|
|
SpeechToText,
|
|
TextChannel,
|
|
TextToSpeech,
|
|
Utterance,
|
|
VisionBackend,
|
|
)
|
|
from .state import SharedScreenContext
|
|
|
|
log = logging.getLogger("wsai.pipeline")
|
|
|
|
|
|
class Pipeline:
|
|
"""Runs two concurrent loops:
|
|
|
|
* perception: FrameSource -> VisionBackend -> SharedScreenContext
|
|
* conversation: (SpeechToText | TextChannel) -> Brain -> (TextToSpeech | TextChannel)
|
|
|
|
Either input/output half can be None, so you can run text-only, voice-only,
|
|
or a headless "just watch" configuration.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
source: FrameSource,
|
|
vision: VisionBackend,
|
|
brain: Brain,
|
|
stt: SpeechToText | None = None,
|
|
tts: TextToSpeech | None = None,
|
|
text_channel: TextChannel | None = None,
|
|
history_turns: int = 12,
|
|
) -> None:
|
|
self.source = source
|
|
self.vision = vision
|
|
self.brain = brain
|
|
self.stt = stt
|
|
self.tts = tts
|
|
self.text_channel = text_channel
|
|
self.context = SharedScreenContext()
|
|
self._history: list[tuple[str, str]] = []
|
|
self._history_turns = history_turns
|
|
|
|
# -- perception -------------------------------------------------------- #
|
|
async def _perceive(self) -> None:
|
|
async for frame in self.source.frames():
|
|
try:
|
|
obs = await self.vision.describe(frame)
|
|
except Exception: # a single bad frame must not kill the loop
|
|
log.exception("vision.describe failed")
|
|
continue
|
|
await self.context.update(obs)
|
|
log.debug("screen: %s", obs.text[:120])
|
|
|
|
# -- conversation ------------------------------------------------------ #
|
|
async def _handle(self, utt: Utterance) -> None:
|
|
screen = await self.context.latest()
|
|
reply = await self.brain.respond(utt.text, screen, self._history)
|
|
self._remember(utt.text, reply.text)
|
|
await self._emit(reply)
|
|
|
|
def _remember(self, user: str, ai: str) -> None:
|
|
self._history.append((user, ai))
|
|
if len(self._history) > self._history_turns:
|
|
self._history = self._history[-self._history_turns :]
|
|
|
|
async def _emit(self, reply: Reply) -> None:
|
|
tasks = []
|
|
if self.tts is not None:
|
|
tasks.append(self.tts.speak(reply))
|
|
if self.text_channel is not None:
|
|
tasks.append(self.text_channel.send(reply))
|
|
if not tasks:
|
|
log.info("AI: %s", reply.text)
|
|
else:
|
|
await asyncio.gather(*tasks)
|
|
|
|
async def _listen_voice(self) -> None:
|
|
if self.stt is None:
|
|
return
|
|
async for utt in self.stt.utterances():
|
|
await self._handle(utt)
|
|
|
|
async def _listen_text(self) -> None:
|
|
if self.text_channel is None:
|
|
return
|
|
async for utt in self.text_channel.messages():
|
|
await self._handle(utt)
|
|
|
|
# -- lifecycle --------------------------------------------------------- #
|
|
async def run(self) -> None:
|
|
loops = [self._perceive(), self._listen_voice(), self._listen_text()]
|
|
try:
|
|
await asyncio.gather(*loops)
|
|
finally:
|
|
await self.aclose()
|
|
|
|
async def aclose(self) -> None:
|
|
for closer in (self.source, self.stt, self.text_channel):
|
|
if closer is not None:
|
|
try:
|
|
await closer.aclose()
|
|
except Exception:
|
|
log.exception("error closing %s", closer)
|