提交: d4e0f10afb105d6ea0b70c43cdb4d3127c60eca0 - zirpu (git)

Emergent generative agents

Commit MetaInfo

• 修訂: d4e0f10afb105d6ea0b70c43cdb4d3127c60eca0 (tree)
• 時間: 2023-05-05 05:11:41
• 作者: Corbin <cds@corb...>
• Commiter: Corbin

Log Message

Switch to Twisted.

Change Summary

• modified: flake.nix
• modified: src/agent.py
• modified: src/common.py

差異

--- a/flake.nix
+++ b/flake.nix
@@ -48,7 +48,8 @@
           doCheck = false;
         };
         py = pkgs.python310.withPackages (ps: with ps; [
-          faiss irc llama-cpp-python sentence-transformers tokenizers transformers torch
+          faiss llama-cpp-python sentence-transformers tokenizers transformers torch
+          twisted
         ]);
         rwkv = pkgs.stdenv.mkDerivation {
           name = "rwkv.cpp";

--- a/src/agent.py
+++ b/src/agent.py
@@ -1,6 +1,5 @@
 #!/usr/bin/env python3
 
-import asyncio
 from collections import deque
 from datetime import datetime
 import json
@@ -8,11 +7,14 @@ import os.path
 import re
 import random
 import sys
-from threading import RLock
+from threading import Lock
 
-from irc.client_aio import AioReactor
-from irc.bot import SingleServerIRCBot
-from irc.strings import lower
+from twisted.internet import reactor
+from twisted.internet.defer import succeed
+from twisted.internet.protocol import ClientFactory
+from twisted.internet.task import LoopingCall
+from twisted.internet.threads import deferToThread
+from twisted.words.protocols.irc import IRCClient
 
 from common import irc_line, Timer, SentenceIndex, breakAt
 # from gens.camelid import CamelidGen
@@ -28,203 +30,204 @@ def load_character(path):
     with open(os.path.join(path, "character.json"), "r") as handle:
         return json.load(handle)
 
-logpath = sys.argv[1]
-character = load_character(logpath)
-startingChannels = character.pop("startingChannels")
-title = character["title"]
-
 MAX_NEW_TOKENS = 128
 print("~ Initializing mawrkov adapter…")
 gen = MawrkovGen(MAX_NEW_TOKENS)
+# Need to protect per-gen data structures in C.
+genLock = Lock()
 GiB = 1024 ** 3
 print("Initialized", gen.model_name, "(" + gen.model_arch + ")",
       "Memory footprint:", round(gen.footprint() / GiB, 2), "GiB")
 embedder = SentenceEmbed()
-thoughtPath = os.path.join(logpath, "thoughts.txt")
-thoughtIndex = SentenceIndex.fromPath(thoughtPath, embedder)
-print("~ Thought index:", thoughtIndex.size(), "thoughts")
+
+prologues = {
+    "clock": "Then I checked the time:",
+    "lbi": "Then I tried to interpret what just happened:",
+    "thoughts": "Then I thought to myself:",
+    "irc": "Then I chatted on IRC:",
+}
 
 class Mind:
-    currentAgent = None
+    currentTag = None
     logits = state = None
 
-    def __init__(self): self.lock = RLock()
+    def switchTag(self, tag):
+        if tag == self.currentTag: return succeed(None)
+        else:
+            self.currentTag = tag
+            # Double newlines are added here.
+            return deferToThread(self.write, "\n" + prologues[tag])
 
     def overhear(self, tag, s):
-        # XXX tag -> highlight agent first!
-        self.write(s)
+        d = self.switchTag(tag)
+        d.addCallback(lambda _: deferToThread(self.write, s))
+        return d
 
     def write(self, s):
-        with self.lock:
+        with genLock:
             print("~ Write:", s)
             # Newlines are added here.
             self.logits, self.state = gen.feedForward(gen.tokenize(s + "\n"),
                                                       self.logits, self.state)
 
     def complete(self, s):
-        with self.lock:
+        with genLock:
             print("~ Completion with prefix:", s)
             completion, self.logits, self.state = gen.complete(s, self.logits, self.state)
             print("~ Completion:", completion)
-            line = s + completion
-            for listener in self.listeners: listener.listenTo(line)
             return completion
 
-    async def infer(self, tag, prefix):
-        # XXX
-        agent = ...
-        # XXX thread all writes and inference
-        self.write("")
-        self.write(agent.prologue())
-        # XXX   for line in agent.activate(): self.write(line)
-        # XXX get all of the messages matching the tag, feed them
-        for message in self.q[tag]:
-            self.write(message)
+    def infer(self, tag, prefix):
         print("~ agent prefix length (tokens):", gen.countTokens(prefix))
-        return self.complete(prefix)
+        d = self.switchTag(tag)
+        d.addCallback(lambda _: deferToThread(self.complete, prefix))
+        return d
 
 
 class Agent:
     listeners = ()
     def broadcast(self, s):
-        for listener in self.listeners:
-            loop.call_soon(listener.overhear, self.tag, s)
+        for listener in self.listeners: listener.overhear(self.tag, s)
 
 class Clock(Agent):
     tag = "clock"
-    def prologue(self): return ["Then I checked the time:"]
-    async def start(self):
-        while True:
-            self.broadcast(f"The time is currently {datetime.now():%H:%M:%S, %B %d %Y}.")
-            asyncio.sleep(60.0)
+    def go(self):
+        self.broadcast(f"The time is currently {datetime.now():%H:%M:%S, %B %d %Y}.")
 
 class LeftBrainInterpreter(Agent):
     tag = "lbi"
     events = 0
-    def prologue(self):
-        return ["Then I tried to interpret what just happened:"]
-    def overhear(self, tag, s): self.events += 1
-
-IRC_LINE_HEAD = re.compile(r"\d{1,2}:\d{1,2}:\d{1,2} <")
-def breakIRCLine(line):
-    return IRC_LINE_HEAD.split(breakAt(line.strip(), "\n"), maxsplit=1)[0]
-
-class IRCAgent(Agent, SingleServerIRCBot):
-    reactor_class = AioReactor
-
-    tag = "irc", None
-    def __init__(self, host, title, startingChannels):
-        super(IRCAgent, self).__init__([(host, 6667)], title_to_nick(title), title)
-        self.startingChannels = startingChannels
-
-    def prologue(self): return ["There is activity in IRC."]
-
-    def highlightChannel(self, channel):
-        if self.highlightedChannel != channel:
-            self.highlightedChannel = channel
-            self.log(f"In channel {channel}:")
-
-    def prefix(self):
-        return f"{datetime.now():%H:%M:%S} <{self.connection.get_nickname()}>"
-
-    def log(self, s):
-        # XXX channel?
-        self.broadcast(s)
-
-    def on_join(self, c, e):
-        who = e.source.nick
-        channel = e.target
-        if who == c.get_nickname(): c.topic(channel)
-        self.log(f"{who} joins {channel}")
-
-    def on_part(self, c, e):
-        who = e.source.nick
-        channel = e.target
-        self.log(f"{who} leaves {channel}")
-
-    def on_currenttopic(self, c, e):
-        channel = e.arguments[0]
-        topic = e.arguments[1]
-        self.channels[channel].topic = topic
-        self.log(f"Topic for {channel} is now: {topic}")
-
-    def on_welcome(self, c, e):
-        for channel in self.startingChannels: c.join(channel)
-
-    def on_pubmsg(self, c, e):
-        line = e.arguments[0]
-        channel = e.target
-        self.highlightChannel(channel)
-        self.log(irc_line(datetime.now(), e.source.nick, line))
-        # Vaguely inspired by
-        # https://github.com/jaraco/irc/blob/main/scripts/testbot.py
-        nick = lower(self.connection.get_nickname())
-        lowered = lower(line)
-        if nick in lowered:
-            print("~ Will respond on IRC")
-            # XXX channel?
-            line = breakIRCLine(self.mind.infer()).strip()
-            self.connection.privmsg(channel, line)
-
+    def __init__(self, mind): self.mind = mind
+    def overhear(self, tag, s):
+        self.events += 1
+        if self.events >= 10:
+            self.events = 0
+            return self.mind.infer(self.tag, "")
 
 class ChainOfThoughts(Agent):
-    def __init__(self, index, seed):
-        super(ChainOfThoughts, self).__init__()
+    tag = "thoughts"
+    def __init__(self, mind, index, seed):
+        self.mind = mind
         self.index = index
         self.recentThoughts = deque([seed], maxlen=5)
-        self.q = []
 
-    async def start(self):
-        while True:
-            cb = self.reflect if random.choice([0, 1]) else self.cogitate
-            await cb()
-            asyncio.sleep(30.0)
-
-    def prologue(self): return ["Thinking to myself:"]
-    def prefix(self): return ""
+    def go(self):
+        cb = self.reflect if random.choice([0, 1]) else self.cogitate
+        return cb()
 
     def addRelatedThoughts(self, s):
-        thoughts = thoughtIndex.search(s, 2)
+        thoughts = self.index.search(s, 2)
         for thought in thoughts:
             if thought not in self.recentThoughts:
                 print("~ New relevant thought:", thought)
                 self.recentThoughts.append(thought)
-                self.q.append(thought)
+                self.broadcast(thought)
 
-    listenTo = addRelatedThoughts
+    def cogitate(self): self.addRelatedThoughts(self.recentThoughts[-1])
+    def reflect(self):
+        d = self.mind.infer(self.tag, "")
 
-    async def cogitate(self):
-        self.addRelatedThoughts(self.recentThoughts[-1])
-        if self.q:
-            for line in self.q: self.broadcast(line)
-            self.q = []
+        @d.addCallback
+        def cb(s):
+            if not s.strip():
+                self.broadcast(random.choice([
+                    "Head empty; no thoughts.",
+                    "So bored.",
+                    "Zoned out.",
+                ]))
 
-    async def reflect(self):
-        # XXX no prefix, just complete
-        self.mind.highlightAgent(self)
-        thought = await self.mind.complete(self.tag, "")
-        print("~ Reflection:", thought)
 
-mind = Mind()
-firstStatement = f"I am {title}."
-with Timer("initial warmup"):
-    mind.logits, mind.state = gen.feedForward(gen.tokenize(firstStatement), None, None)
+IRC_LINE_HEAD = re.compile(r"\d{1,2}:\d{1,2}:\d{1,2} <")
+def breakIRCLine(line):
+    return IRC_LINE_HEAD.split(breakAt(line.strip(), "\n"), maxsplit=1)[0]
 
-agent = IRCAgent("june.local", title, startingChannels)
-loop = agent.reactor.loop
+class IRCAgent(Agent, IRCClient):
+    tag = "irc"
+    def __init__(self, mind, title, startingChannels):
+        super(IRCAgent, self).__init__()
+        self.mind = mind
+        self.nickname = title_to_nick(title)
+        self.startingChannels = startingChannels
 
-clock = Clock()
-lbi = LeftBrainInterpreter()
-thoughts = ChainOfThoughts(thoughtIndex, firstStatement)
-clock.listeners = lbi, mind
-agent.listeners = lbi, mind
-thoughts.listeners = lbi, mind
+    def prefix(self, channel):
+        return f"{datetime.now():%H:%M:%S} {channel} <{self.nickname}>"
+
+    # def on_join(self, c, e):
+    #     who = e.source.nick
+    #     channel = e.target
+    #     if who == c.get_nickname(): c.topic(channel)
+    #     self.broadcast(f"{who} joins {channel}")
+
+    # def on_part(self, c, e):
+    #     who = e.source.nick
+    #     channel = e.target
+    #     self.broadcast(f"{who} leaves {channel}")
+
+    # def on_currenttopic(self, c, e):
+    #     channel = e.arguments[0]
+    #     topic = e.arguments[1]
+    #     self.channels[channel].topic = topic
+    #     self.broadcast(f"Topic for {channel} is now: {topic}")
+
+    def signedOn(self):
+        for channel in self.startingChannels: self.join(channel)
+
+    def privmsg(self, user, channel, line):
+        user = user.split("!", 1)[0]
+        self.broadcast(irc_line(datetime.now(), channel, user, line))
+        if self.nickname in line:
+            print("~ Ping on IRC:", self.nickname)
+            d = self.mind.infer("irc", self.prefix(channel))
+
+            @d.addCallback
+            def cb(s):
+                line = breakIRCLine(s).strip()
+                self.msg(channel, line)
+
+class IRCFactory(ClientFactory):
+    protocol = IRCAgent
+    def __init__(self, mind, listeners, title, startingChannels):
+        super(IRCFactory, self).__init__()
+        self.mind = mind
+        self.listeners = listeners
+        self.title = title
+        self.startingChannels = startingChannels
+    def buildProtocol(self, addr):
+        protocol = self.protocol(self.mind, self.title, self.startingChannels)
+        protocol.factory = self
+        protocol.listeners = self.listeners
+        return protocol
 
 def go():
     print("~ Starting tasks…")
-    asyncio.create_task(clock.start())
-    asyncio.create_task(thoughts.start())
-    agent.start()
+    clock = Clock()
+    LoopingCall(clock.go).start(120.0, now=False)
+
+    for logpath in sys.argv[1:]:
+        character = load_character(logpath)
+        title = character["title"]
+        firstStatement = f"I am {title}."
+        thoughtPath = os.path.join(logpath, "thoughts.txt")
+        thoughtIndex = SentenceIndex.fromPath(thoughtPath, embedder)
+
+        mind = Mind()
+        with Timer("initial warmup"):
+            mind.logits, mind.state = gen.feedForward(gen.tokenize(firstStatement), None, None)
+
+        lbi = LeftBrainInterpreter(mind)
+        clock.listeners += lbi, mind
+
+        thoughts = ChainOfThoughts(mind, thoughtIndex, firstStatement)
+        thoughts.listeners = lbi, mind
+        LoopingCall(thoughts.go).start(60.0, now=False)
+
+        print("~ Thought index:", thoughtIndex.size(), "thoughts")
+        factory = IRCFactory(mind, (lbi, mind),
+                             title,
+                             character["startingChannels"])
+        print("~ Connecting factory for:", title)
+        reactor.connectTCP("june.local", 6667, factory)
+    print("~ Starting event loop…")
+    reactor.run()
 
-print("~ Starting event loop…")
 go()

--- a/src/common.py
+++ b/src/common.py
@@ -1,4 +1,3 @@
-from bisect import bisect
 from itertools import islice
 import random
 from time import perf_counter
@@ -13,36 +12,8 @@ class Timer:
     def __exit__(self, *args):
         print("Timer:", self.l, "%0.02f" % (perf_counter() - self.t), "seconds")
 
-def irc_line(t, speaker, entry): return f"{t:%H:%M:%S} <{speaker}> {entry}"
-
-class Log:
-    "Basic scrollback and context management for conversations."
-    cutoff = 0
-    def __init__(self, l):
-        self.l = l
-        self.stamp = len(self.l)
-
-    def raw(self, line):
-        assert "\n" not in line
-        self.l.append(line)
-        print(self.stamp, ">", line)
-        self.stamp += 1
-
-    def push(self, speaker, entry): self.raw(speaker + ": " + entry)
-    def irc(self, t, speaker, entry): self.raw(irc_line(t, speaker, entry))
-
-    def finishPrompt(self, s, prefix):
-        return self.finishPromptAtCutoff(self.cutoff, s, prefix)
-
-    def finishPromptAtCutoff(self, cutoff, s, prefix):
-        return s + "\n".join(self.l[cutoff:]) + "\n" + prefix
-
-    def undo(self): self.l.pop()
-
-    def bumpCutoff(self, max_context_length, prompt_length, prompt, prefix):
-        def keyfunc(i):
-            return -prompt_length(self.finishPromptAtCutoff(i, prompt, prefix))
-        self.cutoff = bisect(range(len(self.l)), -max_context_length, key=keyfunc)
+def irc_line(t, channel, speaker, entry):
+    return f"{t:%H:%M:%S} {channel} <{speaker}> {entry}"
 
 def parsePygmalion(response, speakers):
     "Pygmalion-specific parser to work around fine-tuned quirks."