From 4eeddc4b1ff24c35902fef1ccce0003ea31d5cac Mon Sep 17 00:00:00 2001 From: claude-owner Date: Sun, 9 Aug 2026 02:16:14 +0900 Subject: [PATCH] feat: scaffold watch-screen AI pipeline (mock-runnable skeleton) 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). --- .gitignore | 5 ++ README.md | 61 ++++++++++++++++ requirements.txt | 17 +++++ tests/test_pipeline.py | 48 +++++++++++++ wsai/__init__.py | 8 +++ wsai/__main__.py | 62 ++++++++++++++++ wsai/backends/__init__.py | 0 wsai/backends/capture_mss.py | 67 +++++++++++++++++ wsai/backends/claude.py | 81 +++++++++++++++++++++ wsai/backends/mock.py | 99 +++++++++++++++++++++++++ wsai/config.py | 53 ++++++++++++++ wsai/factory.py | 70 ++++++++++++++++++ wsai/interfaces.py | 136 +++++++++++++++++++++++++++++++++++ wsai/pipeline.py | 114 +++++++++++++++++++++++++++++ wsai/state.py | 34 +++++++++ 15 files changed, 855 insertions(+) create mode 100644 .gitignore create mode 100644 README.md create mode 100644 requirements.txt create mode 100644 tests/test_pipeline.py create mode 100644 wsai/__init__.py create mode 100644 wsai/__main__.py create mode 100644 wsai/backends/__init__.py create mode 100644 wsai/backends/capture_mss.py create mode 100644 wsai/backends/claude.py create mode 100644 wsai/backends/mock.py create mode 100644 wsai/config.py create mode 100644 wsai/factory.py create mode 100644 wsai/interfaces.py create mode 100644 wsai/pipeline.py create mode 100644 wsai/state.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..4b7d4f5 --- /dev/null +++ b/.gitignore @@ -0,0 +1,5 @@ +__pycache__/ +*.pyc +.pytest_cache/ +.venv/ +*.egg-info/ diff --git a/README.md b/README.md new file mode 100644 index 0000000..ebce254 --- /dev/null +++ b/README.md @@ -0,0 +1,61 @@ +# watch_sceen_ai + +디스코드 화면공유를 **실시간으로 함께 보면서 대화하는 AI**의 골격(skeleton). + +지금 단계는 "틀": 파이프라인 구조와 교체 가능한 인터페이스를 먼저 세우고, +의존성/키/마이크 없이 mock 모드로 전체 흐름이 도는 걸 검증한다. + +## 구조 + +``` +FrameSource ──frames──▶ VisionBackend ──observations──▶ [SharedScreenContext] + │ +STT / TextChannel ──utterances──▶ Brain(LLM) ◀───────────────┘ + │ + ▼ + TextToSpeech / TextChannel +``` + +두 개의 async 루프가 동시에 돈다. +- 인지 루프: 화면 캡처 → 비전 이해 → 최신 화면 맥락 저장 +- 대화 루프: 음성/텍스트 입력 → 두뇌(LLM, 화면맥락+대화이력) → 음성/텍스트 출력 + +모든 단계는 `wsai/interfaces.py`의 Protocol이라, 백엔드(mock / 로컬GPU / 클라우드 / +디스코드 웹캡처 …)를 config에서 바꿔 끼우면 오케스트레이터는 그대로다. + +## 실행 + +```bash +python -m wsai # mock 데모 (의존성·키 불필요, 몇 프레임 돌고 종료) +python -m wsai --live # 이 PC 화면 캡처 + Claude 눈/두뇌 (Ctrl-C 종료) +python -m wsai --env # WSAI_* 환경변수로 조립 +pytest # 스모크 테스트 +``` + +`--live`에 필요한 것: `pip install mss pillow anthropic`, 그리고 `ANTHROPIC_API_KEY`. + +## 파일 + +| 파일 | 역할 | +|------|------| +| `wsai/interfaces.py` | 데이터 타입 + 컴포넌트 Protocol | +| `wsai/pipeline.py` | 오케스트레이터(인지/대화 루프) | +| `wsai/state.py` | 최신 화면 맥락 공유 저장소 | +| `wsai/config.py` / `factory.py` | 설정 → 백엔드 조립 | +| `wsai/backends/mock.py` | 무의존성 mock 전 계열 | +| `wsai/backends/capture_mss.py` | 로컬 화면 캡처(눈) | +| `wsai/backends/claude.py` | Claude 비전 + 두뇌 | + +## 아직 안 된 것 (다음 단계 후보) + +- STT 실제 백엔드(faster-whisper) + 마이크 캡처(sounddevice) +- TTS 실제 백엔드(MeloTTS 등) — 이 호스트에 기존 자산 있음 +- 디스코드 텍스트 채널 I/O (`factory._text`가 아직 NotImplemented) +- 화면 변화 감지 → AI가 먼저 말 거는 proactive 모드 +- 캡처 방식 선택: 로컬 화면 vs 디스코드 웹 캡처(Playwright) + +## 캡처 방식에 대한 메모 + +디스코드 봇이 화면공유 **영상 스트림을 직접 수신**하는 건 공식 API 미지원(유저봇은 ToS 위반). +그래서 현실적 방식은 "통화에 참여해 공유화면을 보는 상태의 화면을 로컬 캡처"하는 것. +`MSSFrameSource`가 그 기본 구현이고, 필요하면 웹 캡처 소스로 교체 가능. diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..3ffd9f1 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,17 @@ +# Core skeleton has NO required third-party deps (mock mode is pure stdlib). +# Install extras per backend you enable: + +# --- screen capture (WSAI_SOURCE=mss) --- +# mss +# pillow + +# --- cloud eyes + brain (WSAI_VISION=claude / WSAI_BRAIN=claude) --- +# anthropic + +# --- planned voice backends (not yet implemented) --- +# faster-whisper # STT +# sounddevice # mic capture +# (a TTS engine, e.g. MeloTTS) + +# --- dev --- +# pytest diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py new file mode 100644 index 0000000..73cddf3 --- /dev/null +++ b/tests/test_pipeline.py @@ -0,0 +1,48 @@ +"""Smoke test: the mock pipeline must run end-to-end and route screen context +into the brain's replies.""" + +import asyncio + +from wsai.backends.mock import ( + MockBrain, + MockFrameSource, + MockSTT, + MockTTS, + MockVision, +) +from wsai.pipeline import Pipeline + + +def test_mock_pipeline_runs_and_replies(capsys): + replies: list[str] = [] + + class CapturingTTS(MockTTS): + async def speak(self, reply): + replies.append(reply.text) + + pipe = Pipeline( + source=MockFrameSource(interval=0.05, limit=3), + vision=MockVision(), + brain=MockBrain(), + stt=MockSTT(script=["화면에 뭐 보여?"], interval=0.1), + tts=CapturingTTS(), + ) + + asyncio.run(asyncio.wait_for(pipe.run(), timeout=5)) + + assert replies, "brain produced no reply" + # The reply must embed the screen observation → context reached the brain. + assert "화면:" in replies[0] + + +def test_history_is_bounded(): + pipe = Pipeline( + source=MockFrameSource(limit=0), + vision=MockVision(), + brain=MockBrain(), + history_turns=3, + ) + for i in range(10): + pipe._remember(f"u{i}", f"a{i}") + assert len(pipe._history) == 3 + assert pipe._history[-1] == ("u9", "a9") diff --git a/wsai/__init__.py b/wsai/__init__.py new file mode 100644 index 0000000..da932ed --- /dev/null +++ b/wsai/__init__.py @@ -0,0 +1,8 @@ +"""watch_screen_ai — an AI that watches a shared screen and talks with you.""" + +from .config import Settings +from .factory import build +from .pipeline import Pipeline + +__all__ = ["Settings", "build", "Pipeline"] +__version__ = "0.0.1" diff --git a/wsai/__main__.py b/wsai/__main__.py new file mode 100644 index 0000000..b49ff6f --- /dev/null +++ b/wsai/__main__.py @@ -0,0 +1,62 @@ +"""Entry point. + + python -m wsai # mock pipeline (no deps, no keys) — runs a demo + python -m wsai --live # capture this screen + Claude eyes/brain + python -m wsai --env # build from WSAI_* environment variables + +The mock run is bounded (a few frames + a scripted conversation) so it exits on +its own; --live/--env run until Ctrl-C. +""" + +from __future__ import annotations + +import argparse +import asyncio +import logging + +from .config import Settings +from .factory import build + + +async def _run(settings: Settings, demo: bool) -> None: + if demo: + # Bounded demo so CI / a quick check terminates. + from .backends.mock import MockFrameSource, MockSTT + from .pipeline import Pipeline + + pipe = build(settings) + pipe.source = MockFrameSource(interval=0.3, limit=4) + if pipe.stt is not None: + pipe.stt = MockSTT(interval=0.4) + await pipe.run() + return + await build(settings).run() + + +def main() -> None: + ap = argparse.ArgumentParser(prog="wsai") + ap.add_argument("--live", action="store_true", help="capture screen + Claude backends") + ap.add_argument("--env", action="store_true", help="build from WSAI_* env vars") + ap.add_argument("-v", "--verbose", action="store_true") + args = ap.parse_args() + + logging.basicConfig( + level=logging.DEBUG if args.verbose else logging.INFO, + format="%(levelname)s %(name)s: %(message)s", + ) + + if args.live: + settings, demo = Settings.live(), False + elif args.env: + settings, demo = Settings.from_env(), False + else: + settings, demo = Settings.mock(), True + + try: + asyncio.run(_run(settings, demo)) + except KeyboardInterrupt: + pass + + +if __name__ == "__main__": + main() diff --git a/wsai/backends/__init__.py b/wsai/backends/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/wsai/backends/capture_mss.py b/wsai/backends/capture_mss.py new file mode 100644 index 0000000..fe01cff --- /dev/null +++ b/wsai/backends/capture_mss.py @@ -0,0 +1,67 @@ +"""Local screen capture via `mss`. + +This is the practical "eye": run this on the machine that is in the Discord call +viewing the shared screen, and it captures that monitor/region. Swap in a +discord-web capture later without touching the pipeline. + +Requires: pip install mss pillow +""" + +from __future__ import annotations + +import asyncio +import io +import time +from typing import AsyncIterator + +from ..interfaces import Frame + + +class MSSFrameSource: + def __init__( + self, + *, + monitor: int = 1, + region: dict | None = None, + interval: float = 1.5, + max_width: int = 1280, + ) -> None: + # `region` overrides `monitor`: {"top":.., "left":.., "width":.., "height":..} + self.monitor = monitor + self.region = region + self.interval = interval + self.max_width = max_width + self._sct = None + + def _ensure(self): + if self._sct is None: + import mss # lazy import so mock mode needs no dependency + + self._sct = mss.mss() + return self._sct + + def _grab_png(self) -> tuple[bytes, int, int]: + from PIL import Image + + sct = self._ensure() + area = self.region or sct.monitors[self.monitor] + shot = sct.grab(area) + img = Image.frombytes("RGB", shot.size, shot.rgb) + if img.width > self.max_width: + h = int(img.height * self.max_width / img.width) + img = img.resize((self.max_width, h)) + buf = io.BytesIO() + img.save(buf, format="PNG") + return buf.getvalue(), img.width, img.height + + async def frames(self) -> AsyncIterator[Frame]: + while True: + # mss is blocking; keep the event loop free. + data, w, h = await asyncio.to_thread(self._grab_png) + yield Frame(data=data, width=w, height=h, ts=time.monotonic(), mime="image/png") + await asyncio.sleep(self.interval) + + async def aclose(self) -> None: + if self._sct is not None: + self._sct.close() + self._sct = None diff --git a/wsai/backends/claude.py b/wsai/backends/claude.py new file mode 100644 index 0000000..8c4a0b7 --- /dev/null +++ b/wsai/backends/claude.py @@ -0,0 +1,81 @@ +"""Cloud brain + vision via the Anthropic (Claude) API. + +Both share one client. Vision sends the frame as a base64 image; the brain is a +plain chat call that receives the latest screen description as context. + +Requires: pip install anthropic +Env: ANTHROPIC_API_KEY +""" + +from __future__ import annotations + +import base64 +import os +import time + +from ..interfaces import Frame, Reply, ScreenObservation + + +def _client(api_key: str | None): + import anthropic + + return anthropic.AsyncAnthropic(api_key=api_key or os.environ.get("ANTHROPIC_API_KEY")) + + +class ClaudeVision: + def __init__(self, *, model: str = "claude-sonnet-4-5", api_key: str | None = None) -> None: + self.model = model + self._client = _client(api_key) + + async def describe(self, frame: Frame, hint: str | None = None) -> ScreenObservation: + prompt = hint or ( + "이건 디스코드 화면공유 캡처야. 지금 화면에서 무슨 일이 벌어지는지 " + "2~3문장으로 한국어로 간결하게 설명해줘. 코드/에러/게임/문서 등 맥락을 짚어줘." + ) + b64 = base64.b64encode(frame.data).decode() + resp = await self._client.messages.create( + model=self.model, + max_tokens=300, + messages=[ + { + "role": "user", + "content": [ + { + "type": "image", + "source": {"type": "base64", "media_type": frame.mime, "data": b64}, + }, + {"type": "text", "text": prompt}, + ], + } + ], + ) + text = "".join(b.text for b in resp.content if b.type == "text") + return ScreenObservation(text=text.strip(), ts=frame.ts) + + +class ClaudeBrain: + SYSTEM = ( + "너는 사용자의 디스코드 화면공유를 실시간으로 함께 보는 AI 파트너야. " + "화면 설명을 참고해 자연스러운 반말/존댓말은 사용자에 맞추고, 짧고 대화하듯 답해. " + "화면을 못 봤으면 솔직히 말해." + ) + + def __init__(self, *, model: str = "claude-sonnet-4-5", api_key: str | None = None) -> None: + self.model = model + self._client = _client(api_key) + + async def respond(self, user_text, screen, history) -> Reply: + msgs = [] + for user, ai in history: + msgs.append({"role": "user", "content": user}) + msgs.append({"role": "assistant", "content": ai}) + screen_note = f"[지금 화면] {screen.text}\n\n" if screen else "[지금 화면] (아직 못 읽음)\n\n" + msgs.append({"role": "user", "content": screen_note + user_text}) + resp = await self._client.messages.create( + model=self.model, + max_tokens=400, + system=self.SYSTEM, + messages=msgs, + ) + text = "".join(b.text for b in resp.content if b.type == "text") + return Reply(text=text.strip(), ts=time.monotonic()) diff --git a/wsai/backends/mock.py b/wsai/backends/mock.py new file mode 100644 index 0000000..1731cae --- /dev/null +++ b/wsai/backends/mock.py @@ -0,0 +1,99 @@ +"""Mock backends. These let the full pipeline run with no GPU, no mic, no API +key — so the skeleton is verifiable and gives every real backend a reference +implementation to match. +""" + +from __future__ import annotations + +import asyncio +import itertools +import time +from typing import AsyncIterator + +from ..interfaces import ( + Frame, + Reply, + ScreenObservation, + Utterance, +) + + +class MockFrameSource: + """Emits tiny synthetic frames on a fixed interval.""" + + def __init__(self, interval: float = 1.0, limit: int | None = None) -> None: + self.interval = interval + self.limit = limit + + async def frames(self) -> AsyncIterator[Frame]: + for i in itertools.count(): + if self.limit is not None and i >= self.limit: + return + yield Frame( + data=b"\x89PNG\r\n\x1a\n", # PNG magic; enough for a stub + width=1280, + height=720, + ts=time.monotonic(), + mime="image/png", + ) + await asyncio.sleep(self.interval) + + async def aclose(self) -> None: # nothing to release + return + + +class MockVision: + """Pretends to read the screen. Cycles through a few canned scenes.""" + + SCENES = [ + "VS Code is open with a Python file; a traceback is visible in the terminal.", + "A browser shows a GitHub pull request diff.", + "A game is running; the player is in a menu screen.", + ] + + def __init__(self) -> None: + self._i = 0 + + async def describe(self, frame: Frame, hint: str | None = None) -> ScreenObservation: + scene = self.SCENES[self._i % len(self.SCENES)] + self._i += 1 + return ScreenObservation(text=scene, ts=frame.ts) + + +class MockSTT: + """Feeds a scripted set of user utterances, then goes quiet.""" + + def __init__(self, script: list[str] | None = None, interval: float = 2.0) -> None: + self.script = script or [ + "지금 화면에 뭐 보여?", + "저 에러 왜 나는 거야?", + "고마워", + ] + self.interval = interval + + async def utterances(self) -> AsyncIterator[Utterance]: + for line in self.script: + await asyncio.sleep(self.interval) + yield Utterance(text=line, ts=time.monotonic(), source="voice") + + async def aclose(self) -> None: + return + + +class MockTTS: + """'Speaks' by printing. Real TTS swaps in here.""" + + async def speak(self, reply: Reply) -> None: + print(f"[TTS] {reply.text}") + + +class MockBrain: + """Echo-style brain that references the current screen, so you can see the + screen context actually reaching the conversation loop.""" + + async def respond(self, user_text, screen, history) -> Reply: + seen = screen.text if screen else "아직 화면을 못 읽었어요" + return Reply( + text=f'(화면: "{seen}") 라고 봤어요. 말씀하신 "{user_text}"에 대해 답하자면… [mock]', + ts=time.monotonic(), + ) diff --git a/wsai/config.py b/wsai/config.py new file mode 100644 index 0000000..f15638a --- /dev/null +++ b/wsai/config.py @@ -0,0 +1,53 @@ +"""Configuration. Each field names a backend; the factory maps names -> classes. + +Defaults are all "mock" so the skeleton runs out of the box. Flip individual +fields (via env or code) as real backends land. + +Env overrides (optional): + WSAI_SOURCE, WSAI_VISION, WSAI_STT, WSAI_TTS, WSAI_BRAIN, WSAI_TEXT + WSAI_CAPTURE_INTERVAL +""" + +from __future__ import annotations + +import os +from dataclasses import dataclass + + +@dataclass +class Settings: + source: str = "mock" # mock | mss + vision: str = "mock" # mock | claude + stt: str | None = "mock" # mock | (whisper) | None + tts: str | None = "mock" # mock | (melo) | None + brain: str = "mock" # mock | claude + text: str | None = None # None | (discord) + + capture_interval: float = 1.5 + anthropic_model: str = "claude-sonnet-4-5" + + @classmethod + def from_env(cls) -> "Settings": + def opt(name: str, default): + v = os.environ.get(name) + return default if v is None else (None if v.lower() == "none" else v) + + return cls( + source=opt("WSAI_SOURCE", "mock"), + vision=opt("WSAI_VISION", "mock"), + stt=opt("WSAI_STT", "mock"), + tts=opt("WSAI_TTS", "mock"), + brain=opt("WSAI_BRAIN", "mock"), + text=opt("WSAI_TEXT", None), + capture_interval=float(os.environ.get("WSAI_CAPTURE_INTERVAL", "1.5")), + ) + + @classmethod + def mock(cls) -> "Settings": + return cls() + + @classmethod + def live(cls) -> "Settings": + """A realistic local config: capture this screen, Claude eyes+brain, + mock voice (until STT/TTS backends are wired).""" + return cls(source="mss", vision="claude", brain="claude", stt="mock", tts="mock") diff --git a/wsai/factory.py b/wsai/factory.py new file mode 100644 index 0000000..119e0bd --- /dev/null +++ b/wsai/factory.py @@ -0,0 +1,70 @@ +"""Build a Pipeline from Settings. This is the single place that knows which +concrete class each config name maps to, so adding a backend = one line here.""" + +from __future__ import annotations + +from .config import Settings +from .pipeline import Pipeline + + +def build(settings: Settings) -> Pipeline: + return Pipeline( + source=_source(settings), + vision=_vision(settings), + brain=_brain(settings), + stt=_stt(settings), + tts=_tts(settings), + text_channel=_text(settings), + ) + + +def _source(s: Settings): + if s.source == "mss": + from .backends.capture_mss import MSSFrameSource + + return MSSFrameSource(interval=s.capture_interval) + from .backends.mock import MockFrameSource + + return MockFrameSource(interval=s.capture_interval) + + +def _vision(s: Settings): + if s.vision == "claude": + from .backends.claude import ClaudeVision + + return ClaudeVision(model=s.anthropic_model) + from .backends.mock import MockVision + + return MockVision() + + +def _brain(s: Settings): + if s.brain == "claude": + from .backends.claude import ClaudeBrain + + return ClaudeBrain(model=s.anthropic_model) + from .backends.mock import MockBrain + + return MockBrain() + + +def _stt(s: Settings): + if s.stt in (None, "none"): + return None + from .backends.mock import MockSTT + + return MockSTT() + + +def _tts(s: Settings): + if s.tts in (None, "none"): + return None + from .backends.mock import MockTTS + + return MockTTS() + + +def _text(s: Settings): + if s.text in (None, "none"): + return None + raise NotImplementedError("discord text channel backend not implemented yet") diff --git a/wsai/interfaces.py b/wsai/interfaces.py new file mode 100644 index 0000000..0be55b0 --- /dev/null +++ b/wsai/interfaces.py @@ -0,0 +1,136 @@ +"""Core data types and component interfaces for the watch-screen AI. + +The whole system is a small pipeline: + + FrameSource --frames--> VisionBackend --observations--> [SharedScreenContext] + | + SpeechToText / TextInput --utterances--> Brain <----------------/ + | + v + TextToSpeech / TextOutput + +Every stage is a Protocol so a concrete backend (mock, local GPU, cloud API, +discord web capture, ...) can be swapped in from config without touching the +orchestrator. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import AsyncIterator, Protocol, runtime_checkable + + +# --------------------------------------------------------------------------- # +# Data that flows through the pipeline +# --------------------------------------------------------------------------- # + +@dataclass +class Frame: + """A single captured image of the shared screen.""" + + # Raw encoded image bytes (PNG/JPEG). Kept as bytes so any backend can + # decode it however it likes and so it is trivial to base64 for a cloud API. + data: bytes + width: int + height: int + # Monotonic capture timestamp in seconds. + ts: float + mime: str = "image/png" + + +@dataclass +class ScreenObservation: + """What the vision backend understood from a Frame.""" + + text: str + ts: float + # Optional structured hints (e.g. detected app, code language, error text). + tags: dict[str, str] = field(default_factory=dict) + + +@dataclass +class Utterance: + """Something the user said (voice→text) or typed.""" + + text: str + ts: float + source: str = "voice" # "voice" | "text" + + +@dataclass +class Reply: + """The AI's response, ready to be spoken and/or shown.""" + + text: str + ts: float + + +# --------------------------------------------------------------------------- # +# Component interfaces +# --------------------------------------------------------------------------- # + +@runtime_checkable +class FrameSource(Protocol): + """Produces frames of the shared screen.""" + + async def frames(self) -> AsyncIterator[Frame]: + """Yield frames until cancelled. Cadence is up to the implementation.""" + ... + + async def aclose(self) -> None: + ... + + +@runtime_checkable +class VisionBackend(Protocol): + """Turns a Frame into a text description of what is on screen.""" + + async def describe(self, frame: Frame, hint: str | None = None) -> ScreenObservation: + ... + + +@runtime_checkable +class SpeechToText(Protocol): + """Streams user utterances from the microphone (or a mock source).""" + + async def utterances(self) -> AsyncIterator[Utterance]: + ... + + async def aclose(self) -> None: + ... + + +@runtime_checkable +class TextToSpeech(Protocol): + """Speaks a reply out loud.""" + + async def speak(self, reply: Reply) -> None: + ... + + +@runtime_checkable +class Brain(Protocol): + """The conversational LLM. Given the latest screen context, the user's + message and the running history, produce a reply.""" + + async def respond( + self, + user_text: str, + screen: ScreenObservation | None, + history: list[tuple[str, str]], + ) -> Reply: + ... + + +@runtime_checkable +class TextChannel(Protocol): + """Optional text I/O (e.g. a Discord channel) that mirrors the voice loop.""" + + async def messages(self) -> AsyncIterator[Utterance]: + ... + + async def send(self, reply: Reply) -> None: + ... + + async def aclose(self) -> None: + ... diff --git a/wsai/pipeline.py b/wsai/pipeline.py new file mode 100644 index 0000000..2042b3e --- /dev/null +++ b/wsai/pipeline.py @@ -0,0 +1,114 @@ +"""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) diff --git a/wsai/state.py b/wsai/state.py new file mode 100644 index 0000000..8c13fb5 --- /dev/null +++ b/wsai/state.py @@ -0,0 +1,34 @@ +"""Shared, thread/async-safe screen context. + +The perception loop keeps writing the latest ScreenObservation here; the +conversation loop reads it when the user says something. We only keep the most +recent observation plus a short ring buffer of recent ones so the Brain can +notice "the screen changed" without us re-sending every frame. +""" + +from __future__ import annotations + +import asyncio +from collections import deque + +from .interfaces import ScreenObservation + + +class SharedScreenContext: + def __init__(self, history: int = 8) -> None: + self._latest: ScreenObservation | None = None + self._recent: deque[ScreenObservation] = deque(maxlen=history) + self._lock = asyncio.Lock() + + async def update(self, obs: ScreenObservation) -> None: + async with self._lock: + self._latest = obs + self._recent.append(obs) + + async def latest(self) -> ScreenObservation | None: + async with self._lock: + return self._latest + + async def recent(self) -> list[ScreenObservation]: + async with self._lock: + return list(self._recent)