okami.mom/albedo

git clone https://git.nkp.pet/okami.mom/albedo

Aly Raffaufrefactor(tests): share repeated E2E fixtures5bdc0ee

main
branchesmainall refs →
8.6 KiB215 linesraw
1"""Python agent API: spawn, family identity, refusals, progress, explicit mail,
2agents.get, reading messages across sessions, and listing and searching them."""
3
4import re
5import time
6import unittest
7
8from harness import Albedo, Provider, latest_user, python, text
9
10LEAD = """models = await agents.models()
11assert 'beta/fixture-beta' in models, models
12kid = await agents.self.spawn('count to three', name='scout', model=models[0])
13try:
14    await agents.self.spawn('again', name='scout2', model='')
15except TypeError as error:
16    print('EMPTY_MODEL', error)
17print('SPAWNED', kid.name, kid.depth, kid.parent.name, 'KID=' + kid.id)
18"""
19CHILD = """me = agents.self
20print('ME', me.name, me.depth, me.parent.name)
21print('PROGRESS', await agents.progress('counting'))
22print('Café Ünïcode')
23try:
24    await me.parent.cancel()
25except AgentsError as error:
26    print('REFUSED', error)
27page = await me.parent.messages()
28print('PARENT_READ', 'spawn a scout' in page.content)
29up = await agents.get('parent')
30print('GOT', up.id == me.parent.id, up.depth, up.parent, up.closed)
31try:
32    await me.parent.delete()
33except AgentsError as error:
34    print('ASK', error)
35receipt = await mail.submit('parent', 'three')
36print('SENT', receipt.status, receipt.name)
37"""
38
39READ = """kid = (await agents.self.children())[0]
40page = await kid.messages()
41print('TASK_SEEN', 'count to three' in page.content)
42hits = await kid.search_messages('COUNT TO THREE')
43seq = hits.rows[0]['seq']
44row = await kid.messages(seq=seq, limit=200)
45print('FOUND', hits.count, row.content.startswith(f'[row #{seq}]'), 'count to three' in row.content)
46"""
47
48PEEK = """async def paged_sessions(query='', cwd=None):
49    found = []
50    while True:
51        page = await agents.sessions(query, cwd=cwd, offset=len(found))
52        if not page:
53            return found
54        found.extend(page)
55
56kid = await agents.get({kid!r})
57print('PEEK', kid.name, kid.depth, kid.closed)
58page = await kid.messages()
59print('PEEK_READ', 'count to three' in page.content)
60try:
61    await kid.cancel()
62except AgentsError as error:
63    print('PEEK_REFUSED', error)
64listed = await paged_sessions()
65print('LISTED', {kid!r} in [s.id for s in listed], agents.self.id in [s.id for s in listed])
66child_session = next(s for s in listed if s.id == {kid!r})
67scouts = await paged_sessions('SCOUT', cwd=child_session.cwd)
68scout = next(s for s in scouts if s.id == {kid!r})
69print('BY_NAME', scout.name, scout.depth, scout.model, scout.cwd == child_session.cwd)
70print('ELSEWHERE', await paged_sessions(cwd='/nowhere'))
71talked = await paged_sessions('COUNT TO THREE')
72said = next(s for s in talked if s.id == {kid!r})
73hit = said.matches[0]
74row = await said.messages(seq=hit['seq'], limit=200)
75print('SEARCHED', agents.self.id in [s.id for s in talked], 'count to three' in hit['preview'], 'count to three' in row.content)
76print('UNQUERIED', listed[0].matches)
77print('UNICODE', {kid!r} in [s.id for s in await paged_sessions('CAFÉ ÜNÏCODE')])
78print('NOTHING', await paged_sessions('never ' + 'said anywhere'))
79try:
80    await agents.get('nobody')
81except AgentsError as error:
82    print('PEEK_MISSING', error)
83"""
84
85
86def wait_for(predicate):
87    deadline = time.monotonic() + 60
88    while time.monotonic() < deadline:
89        result = predicate()
90        if result:
91            return result
92        time.sleep(0.1)
93    raise AssertionError("agent Python turn never arrived")
94
95
96class AgentsPythonTests(unittest.TestCase):
97    def test_python_spawn_identity_refusals_and_explicit_answer(self):
98        def script(request):
99            last = request["messages"][-1]
100            user = latest_user(request)
101            peek = re.search(r"peek ([0-9a-f]+)", user)
102            if last.get("role") == "user" and peek:
103                return python(PEEK.format(kid=peek.group(1)))
104            if last.get("role") == "user" and user == "spawn a scout":
105                return python(LEAD)
106            if last.get("role") == "user" and 'kind="task"' in user:
107                return python(CHILD)
108            if last.get("role") == "user" and 'kind="message"' in user:
109                return python(READ)
110            return text("ok")
111
112        provider = Provider(script)
113        providers = {
114            name: {
115                "baseUrl": provider.url + f"/{name}/v1",
116                "apiKey": "key",
117                "model": f"fixture-{name}",
118                "protocol": "chat_completions",
119            }
120            for name in ("alpha", "beta")
121        }
122        try:
123            with Albedo(provider, providers=providers) as app:
124                lead = app.session()
125                app.prompt(lead, "spawn a scout").close()
126
127                def observed():
128                    records = []
129                    for item in provider.requests:
130                        messages = item["request"]["messages"]
131                        last = messages[-1]
132                        records.append(
133                            {
134                                "user": latest_user(item["request"]),
135                                "context": " ".join(
136                                    str(m.get("content", "")) for m in messages
137                                ),
138                                "tool": str(last.get("content", ""))
139                                if last.get("role") == "tool"
140                                else "",
141                            }
142                        )
143                    return records
144
145                spawned = wait_for(
146                    lambda: next(
147                        (r["tool"] for r in observed() if "SPAWNED" in r["tool"]), None
148                    )
149                )
150                self.assertIn("SPAWNED scout 1", spawned)
151                self.assertIn("EMPTY_MODEL", spawned)
152                child = wait_for(
153                    lambda: next(
154                        (r["tool"] for r in observed() if "SENT" in r["tool"]), None
155                    )
156                )
157                self.assertIn("ME scout 1", child)
158                self.assertIn("PROGRESS True", child)
159                self.assertIn("REFUSED", child)
160                self.assertIn("own children", child)
161                self.assertIn("PARENT_READ True", child)
162                self.assertIn("GOT True 0 None False", child)
163                self.assertIn("ASK", child)
164                self.assertIn("only the user deletes", child)
165                self.assertTrue(
166                    "SENT delivered spawn a scout" in child
167                    or "SENT queued spawn a scout" in child
168                )
169                task = next(r for r in observed() if 'kind="task"' in r["user"])
170                self.assertIn('You are child agent "scout"', task["context"])
171                wait_for(
172                    lambda: any(
173                        'kind="message"' in r["user"] and "three" in r["user"]
174                        for r in observed()
175                    )
176                )
177                read = wait_for(
178                    lambda: next(
179                        (r["tool"] for r in observed() if "FOUND" in r["tool"]), None
180                    )
181                )
182                self.assertIn("TASK_SEEN True", read)
183                self.assertRegex(read, r"FOUND [1-9]\d* True True")
184                match = re.search(r"KID=([0-9a-f]+)", spawned)
185                assert match is not None
186                kid = match.group(1)
187                outsider = app.session()
188                app.prompt(outsider, f"peek {kid}").close()
189                app.idle(outsider, timeout=60)
190                delivered = [record["tool"] for record in observed() if record["tool"]]
191                peeked = next(
192                    (record for record in delivered if "PEEK_MISSING" in record), None
193                )
194                assert peeked is not None, (
195                    "PEEK turn finished without its final marker; delivered tools: "
196                    + repr(delivered)
197                )
198                self.assertIn("PEEK scout 1 False", peeked)
199                self.assertIn("PEEK_READ True", peeked)
200                self.assertIn("PEEK_REFUSED you can only cancel or close", peeked)
201                self.assertIn("no agent named 'nobody'", peeked)
202                self.assertIn("LISTED True True", peeked)
203                self.assertIn("BY_NAME scout 1 fixture-alpha True", peeked)
204                self.assertIn("ELSEWHERE []", peeked)
205                self.assertIn("SEARCHED False True True", peeked)
206                self.assertIn("UNQUERIED []", peeked)
207                self.assertIn("UNICODE True", peeked)
208                self.assertIn("NOTHING []", peeked)
209                app.idle(lead)
210                time.sleep(1)
211                self.assertFalse(
212                    any('kind="unreviewed"' in r["user"] for r in observed())
213                )
214        finally:
215            provider.close()