discli

Time:  0s      1.5s     3.0s     4.5s     5.0s
       │        │        │        │        │
Chunks: Hello,  I can    help     with     that!
Edits:  "..."   "Hello, I can "  "Hello, I can help with "  (stream_end) "Hello, I can help with that!"

import json
import subprocess
import threading
import time
from queue import Queue


class StreamingAgent:
    def __init__(self):
        self.proc = subprocess.Popen(
            ["discli", "serve", "--events", "messages"],
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            text=True,
        )
        self.events = Queue()
        self.pending = {}
        self._counter = 0

        threading.Thread(target=self._reader, daemon=True).start()

    def _reader(self):
        for line in self.proc.stdout:
            data = json.loads(line.strip())
            rid = data.get("req_id")
            if rid in self.pending:
                self.pending[rid].put(data)
            else:
                self.events.put(data)

    def send(self, action, wait=False):
        self._counter += 1
        rid = f"r{self._counter}"
        action["req_id"] = rid
        if wait:
            q = Queue()
            self.pending[rid] = q
        self.proc.stdin.write(json.dumps(action) + "\n")
        self.proc.stdin.flush()
        if wait:
            result = q.get(timeout=30)
            del self.pending[rid]
            return result
        return None

    def stream_to_discord(self, channel_id, reply_to, text_generator):
        """將生成器輸出的文字串流到 Discord 訊息。"""

        # 步驟一：啟動串流
        self.send({"action": "typing_start", "channel_id": channel_id})
        result = self.send({
            "action": "stream_start",
            "channel_id": channel_id,
            "reply_to": reply_to,
        }, wait=True)
        self.send({"action": "typing_stop", "channel_id": channel_id})

        if "error" in result:
            # 備援：若串流失敗，將內容整段收集後直接發送
            full_text = "".join(text_generator)
            self.send({
                "action": "reply",
                "channel_id": channel_id,
                "message_id": reply_to,
                "content": full_text,
            })
            return

        stream_id = result["stream_id"]

        # 步驟二：逐片段寫入
        try:
            for chunk in text_generator:
                self.send({
                    "action": "stream_chunk",
                    "stream_id": stream_id,
                    "content": chunk,
                })
        except Exception as e:
            # 以錯誤訊息作為最後片段
            self.send({
                "action": "stream_chunk",
                "stream_id": stream_id,
                "content": f"\n\n[Error: {e}]",
            })

        # 步驟三：結束
        self.send({"action": "stream_end", "stream_id": stream_id}, wait=True)


def call_llm(prompt):
    """呼叫 OpenAI 相容 API，並逐段回傳回覆內容。"""
    import openai

    client = openai.OpenAI()  # 使用 OPENAI_API_KEY 環境變數
    stream = client.chat.completions.create(
        model="gpt-4o",
        messages=[{"role": "user", "content": prompt}],
        stream=True,
    )
    for chunk in stream:
        delta = chunk.choices[0].delta.content
        if delta:
            yield delta


def main():
    agent = StreamingAgent()

    # 等待 ready
    while True:
        event = agent.events.get()
        if event.get("event") == "ready":
            print(f"Ready as {event['bot_name']}")
            break

    # 處理訊息
    while True:
        event = agent.events.get()
        if event.get("event") != "message":
            continue
        if not event.get("mentions_bot") or event.get("is_bot"):
            continue

        agent.stream_to_discord(
            channel_id=event["channel_id"],
            reply_to=event["message_id"],
            text_generator=call_llm(event["content"]),
        )


if __name__ == "__main__":
    main()

def handle_slash_command(agent, event):
    stream_result = agent.send({
        "action": "stream_start",
        "channel_id": event["channel_id"],
        "interaction_token": event["interaction_token"],
    }, wait=True)

    stream_id = stream_result["stream_id"]

    for chunk in call_llm(event["args"].get("question", "")):
        agent.send({
            "action": "stream_chunk",
            "stream_id": stream_id,
            "content": chunk,
        })

    agent.send({"action": "stream_end", "stream_id": stream_id})