Skip to content

Commit 0002d64

Browse files
jph00jph00
committed
fixes #226 
1 parent cba6f13 commit 0002d64

4 files changed

Lines changed: 181 additions & 2 deletions

File tree

‎dialoghelper/_modidx.py‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,15 @@
1010
'dialoghelper.capture.capture_tool': ('capture.html#capture_tool', 'dialoghelper/capture.py'),
1111
'dialoghelper.capture.setup_share': ('capture.html#setup_share', 'dialoghelper/capture.py'),
1212
'dialoghelper.capture.start_share': ('capture.html#start_share', 'dialoghelper/capture.py')},
13-
'dialoghelper.core': { 'dialoghelper.core._add_msg': ('core.html#_add_msg', 'dialoghelper/core.py'),
13+
'dialoghelper.core': { 'dialoghelper.core.Channel': ('core.html#channel', 'dialoghelper/core.py'),
14+
'dialoghelper.core.Channel.__init__': ('core.html#channel.__init__', 'dialoghelper/core.py'),
15+
'dialoghelper.core.Channel._read_loop': ('core.html#channel._read_loop', 'dialoghelper/core.py'),
16+
'dialoghelper.core.Channel.close': ('core.html#channel.close', 'dialoghelper/core.py'),
17+
'dialoghelper.core.Channel.connect': ('core.html#channel.connect', 'dialoghelper/core.py'),
18+
'dialoghelper.core.Channel.is_open': ('core.html#channel.is_open', 'dialoghelper/core.py'),
19+
'dialoghelper.core.Channel.request': ('core.html#channel.request', 'dialoghelper/core.py'),
20+
'dialoghelper.core.Channel.send': ('core.html#channel.send', 'dialoghelper/core.py'),
21+
'dialoghelper.core._add_msg': ('core.html#_add_msg', 'dialoghelper/core.py'),
1422
'dialoghelper.core._add_msg_unsafe': ('core.html#_add_msg_unsafe', 'dialoghelper/core.py'),
1523
'dialoghelper.core._check_res': ('core.html#_check_res', 'dialoghelper/core.py'),
1624
'dialoghelper.core._diff_dialog': ('core.html#_diff_dialog', 'dialoghelper/core.py'),

‎dialoghelper/core.py‎

Lines changed: 45 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,14 +8,15 @@
88
'call_endp', 'call_endpa', 'curr_dialog', 'msg_idx', 'add_html_a', 'add_html', 'add_scr_a', 'add_scr',
99
'iife_a', 'iife', 'add_mod', 'add_mod_a', 'pop_data_a', 'pop_data', 'fire_event_a', 'fire_event',
1010
'event_get_a', 'event_get', 'trigger_now', 'event_once', 'event_once_a', 'js_run', 'js_run_a', 'js_eval',
11-
'js_eval_a', 'display_response', 'connfiles', 'realpath', 'list_dialogs', 'read_msg', 'find_msgs',
11+
'js_eval_a', 'Channel', 'display_response', 'connfiles', 'realpath', 'list_dialogs', 'read_msg', 'find_msgs',
1212
'view_dlg', 'add_msg', 'read_msgid', 'view_msg', 'msg_ref', 'del_msg', 'run_and_prompt', 'update_msg',
1313
'run_msg', 'copy_msg', 'paste_msg', 'enable_mermaid', 'mermaid', 'toggle_header', 'toggle_bookmark',
1414
'toggle_comment', 'url2note', 'create_or_run_dialog', 'stop_dialog', 'load_dialog', 'rm_dialog',
1515
'run_code_interactive', 'solveit_docs', 'dialog_link', 'spawn_agent', 'search', 'searches', 'web_answer']
1616

1717
# %% ../nbs/00_core.ipynb #4dd4b925
1818
import os,re,inspect,ast,collections,time,asyncio,json,linecache,importlib,difflib,uuid,builtins,subprocess,sys
19+
import websockets
1920

2021
from typing import Dict
2122
from tempfile import TemporaryDirectory
@@ -284,6 +285,49 @@ async def js_eval_a(expr):
284285
await iife_a("pushData('%s', { result: await (async () => { %s })() })" % (idx, expr))
285286
return await pop_data_a(idx)
286287

288+
# %% ../nbs/00_core.ipynb #84469f4e
289+
class Channel:
290+
"Duplex JSON messaging with the other peers on a named `/wsx` relay channel"
291+
def __init__(self, chan, url=None):
292+
self.url = url or f"ws://localhost:{dh_settings['port']}/wsx?chan={chan}"
293+
self._replies,self.events = {},asyncio.Queue()
294+
295+
@classmethod
296+
async def connect(cls,
297+
chan, # Channel name; peers connecting to the same name hear each other
298+
url=None # Full relay url (defaults to local solveit's `/wsx`)
299+
):
300+
"Connect to `chan` and start reading frames"
301+
self = cls(chan, url=url)
302+
self._ws = await websockets.connect(self.url)
303+
self._reader = asyncio.create_task(self._read_loop())
304+
return self
305+
306+
async def _read_loop(self):
307+
async for frame in self._ws:
308+
msg = dict2obj(json.loads(frame))
309+
if (f := self._replies.pop(msg.get('id'), None)): f.set_result(msg)
310+
else: await self.events.put(msg)
311+
312+
async def send(self, **msg): await self._ws.send(json.dumps(msg))
313+
314+
async def request(self, timeout=15, **msg):
315+
"Send `msg` with a correlation `id` and await the reply frame bearing that id"
316+
mid = msg['id'] = str(uuid.uuid4())
317+
f = asyncio.get_running_loop().create_future()
318+
self._replies[mid] = f
319+
try:
320+
await self._ws.send(json.dumps(msg))
321+
return await asyncio.wait_for(f, timeout)
322+
finally: self._replies.pop(mid, None)
323+
324+
async def close(self):
325+
self._reader.cancel()
326+
await self._ws.close()
327+
328+
@property
329+
def is_open(self): return self._ws.state.name == 'OPEN'
330+
287331
# %% ../nbs/00_core.ipynb #80334098
288332
def display_response(display:str, result:str=None):
289333
"Return a special response where `display` is added as markdown/HTML to the prompt output, and `result` is returned to the LLM"

‎nbs/00_core.ipynb‎

Lines changed: 126 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@
3737
"source": [
3838
"#| export\n",
3939
"import os,re,inspect,ast,collections,time,asyncio,json,linecache,importlib,difflib,uuid,builtins,subprocess,sys\n",
40+
"import websockets\n",
4041
"\n",
4142
"from typing import Dict\n",
4243
"from tempfile import TemporaryDirectory\n",
@@ -933,6 +934,131 @@
933934
"await js_eval_a(\"return (await fetch('/curr_dialog_', {method:'POST'})).json()\")"
934935
]
935936
},
937+
{
938+
"cell_type": "markdown",
939+
"id": "fbfeec09",
940+
"metadata": {},
941+
"source": [
942+
"The helpers above make one-shot round trips through the browser page, and remain the right choice for collecting single values. `Channel` instead holds open a persistent two-way connection, for long-lived protocols and event streams where per-message HTTP overhead would hurt."
943+
]
944+
},
945+
{
946+
"cell_type": "code",
947+
"execution_count": null,
948+
"id": "84469f4e",
949+
"metadata": {},
950+
"outputs": [],
951+
"source": [
952+
"#| export\n",
953+
"class Channel:\n",
954+
" \"Duplex JSON messaging with the other peers on a named `/wsx` relay channel\"\n",
955+
" def __init__(self, chan, url=None):\n",
956+
" self.url = url or f\"ws://localhost:{dh_settings['port']}/wsx?chan={chan}\"\n",
957+
" self._replies,self.events = {},asyncio.Queue()\n",
958+
"\n",
959+
" @classmethod\n",
960+
" async def connect(cls,\n",
961+
" chan, # Channel name; peers connecting to the same name hear each other\n",
962+
" url=None # Full relay url (defaults to local solveit's `/wsx`)\n",
963+
" ):\n",
964+
" \"Connect to `chan` and start reading frames\"\n",
965+
" self = cls(chan, url=url)\n",
966+
" self._ws = await websockets.connect(self.url)\n",
967+
" self._reader = asyncio.create_task(self._read_loop())\n",
968+
" return self\n",
969+
"\n",
970+
" async def _read_loop(self):\n",
971+
" async for frame in self._ws:\n",
972+
" msg = dict2obj(json.loads(frame))\n",
973+
" if (f := self._replies.pop(msg.get('id'), None)): f.set_result(msg)\n",
974+
" else: await self.events.put(msg)\n",
975+
"\n",
976+
" async def send(self, **msg): await self._ws.send(json.dumps(msg))\n",
977+
"\n",
978+
" async def request(self, timeout=15, **msg):\n",
979+
" \"Send `msg` with a correlation `id` and await the reply frame bearing that id\"\n",
980+
" mid = msg['id'] = str(uuid.uuid4())\n",
981+
" f = asyncio.get_running_loop().create_future()\n",
982+
" self._replies[mid] = f\n",
983+
" try:\n",
984+
" await self._ws.send(json.dumps(msg))\n",
985+
" return await asyncio.wait_for(f, timeout)\n",
986+
" finally: self._replies.pop(mid, None)\n",
987+
"\n",
988+
" async def close(self):\n",
989+
" self._reader.cancel()\n",
990+
" await self._ws.close()\n",
991+
"\n",
992+
" @property\n",
993+
" def is_open(self): return self._ws.state.name == 'OPEN'"
994+
]
995+
},
996+
{
997+
"cell_type": "markdown",
998+
"id": "4f6edf81",
999+
"metadata": {},
1000+
"source": [
1001+
"A `Channel` connects to a named channel on solveit's `/wsx` relay, which forwards every frame verbatim to all other peers on the same name. Frames are JSON dicts: pass fields as kwargs to `send`, and receive them as `AttrDict`s. Frames that answer a pending `request` are matched up by their `id`; everything else goes to the `events` queue. The main peer today is the [solveit-chrome](https://github.com/AnswerDotAI/solveit-chrome) extension, which `solvecdp` drives through this class.\n",
1002+
"\n",
1003+
"Two peers on one channel: a plain `send` arrives in the other peer's `events`."
1004+
]
1005+
},
1006+
{
1007+
"cell_type": "code",
1008+
"execution_count": null,
1009+
"id": "5c5afd2b",
1010+
"metadata": {},
1011+
"outputs": [],
1012+
"source": [
1013+
"ch1,ch2 = await Channel.connect('test'),await Channel.connect('test')\n",
1014+
"await ch1.send(hello='world')\n",
1015+
"test_eq((await ch2.events.get()).hello, 'world')"
1016+
]
1017+
},
1018+
{
1019+
"cell_type": "markdown",
1020+
"id": "44ae366c",
1021+
"metadata": {},
1022+
"source": [
1023+
"`request` adds an `id` to the frame and waits for the reply with the same `id`, so several requests can be in flight at once. Here the second peer answers, echoing the `id` back with a result:"
1024+
]
1025+
},
1026+
{
1027+
"cell_type": "code",
1028+
"execution_count": null,
1029+
"id": "6e191ec9",
1030+
"metadata": {},
1031+
"outputs": [],
1032+
"source": [
1033+
"async def _responder():\n",
1034+
" msg = await ch2.events.get()\n",
1035+
" await ch2.send(id=msg.id, result=msg.x*2)\n",
1036+
"\n",
1037+
"task = asyncio.create_task(_responder())\n",
1038+
"test_eq((await ch1.request(x=21)).result, 42)"
1039+
]
1040+
},
1041+
{
1042+
"cell_type": "markdown",
1043+
"id": "8354e9fe",
1044+
"metadata": {},
1045+
"source": [
1046+
"`is_open` shows whether the connection is still up:"
1047+
]
1048+
},
1049+
{
1050+
"cell_type": "code",
1051+
"execution_count": null,
1052+
"id": "f24d7639",
1053+
"metadata": {},
1054+
"outputs": [],
1055+
"source": [
1056+
"assert ch1.is_open\n",
1057+
"await ch1.close()\n",
1058+
"await ch2.close()\n",
1059+
"assert not ch1.is_open"
1060+
]
1061+
},
9361062
{
9371063
"cell_type": "code",
9381064
"execution_count": null,

‎pyproject.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ dependencies = [
2626
"safepyrun>=0.0.23",
2727
'ast-grep-cli', 'ast-grep-py',
2828
'restrictedpython', "restrictedpython-async",
29+
'websockets',
2930
'pillow', "pyjwt", "tracefunc"
3031
]
3132

0 commit comments

Comments
 (0)