From 2e5c703603e212a35df421b751fd1dfa78631c3b Mon Sep 17 00:00:00 2001 From: Jeremy Howard Date: Thu, 3 Sep 2026 05:18:40 +1000 Subject: [PATCH] Carry a provider's raw message on `Msg`, keep server tool calls as ordinary pairs, and count PDF pages without pypdf --- aidialog/_modidx.py | 4 +--- aidialog/hist.py | 15 +++++++++---- aidialog/msg_parts.py | 29 ++++++++++-------------- nbs/00_msg_parts.ipynb | 47 +++++++++++++++++++++++++++------------ nbs/03_hist.ipynb | 50 +++++++++++++++++++++++++++++++++++++++++- 5 files changed, 106 insertions(+), 39 deletions(-) diff --git a/aidialog/_modidx.py b/aidialog/_modidx.py index ec907bf..87f5f5b 100644 --- a/aidialog/_modidx.py +++ b/aidialog/_modidx.py @@ -208,6 +208,7 @@ 'aidialog.hist.media_item': ('hist.html#media_item', 'aidialog/hist.py'), 'aidialog.hist.merge_media': ('hist.html#merge_media', 'aidialog/hist.py'), 'aidialog.hist.output_parts': ('hist.html#output_parts', 'aidialog/hist.py'), + 'aidialog.hist.pdf_pages': ('hist.html#pdf_pages', 'aidialog/hist.py'), 'aidialog.hist.reply2chat': ('hist.html#reply2chat', 'aidialog/hist.py'), 'aidialog.hist.reply2dlg': ('hist.html#reply2dlg', 'aidialog/hist.py'), 'aidialog.hist.resize_img': ('hist.html#resize_img', 'aidialog/hist.py'), @@ -275,9 +276,6 @@ 'aidialog.msg_parts.Part.replace': ('msg_parts.html#part.replace', 'aidialog/msg_parts.py'), 'aidialog.msg_parts.Refusal': ('msg_parts.html#refusal', 'aidialog/msg_parts.py'), 'aidialog.msg_parts.Refusal.__init__': ('msg_parts.html#refusal.__init__', 'aidialog/msg_parts.py'), - 'aidialog.msg_parts.ServerToolResult': ('msg_parts.html#servertoolresult', 'aidialog/msg_parts.py'), - 'aidialog.msg_parts.ServerToolResult.__init__': ( 'msg_parts.html#servertoolresult.__init__', - 'aidialog/msg_parts.py'), 'aidialog.msg_parts.StopResponse': ('msg_parts.html#stopresponse', 'aidialog/msg_parts.py'), 'aidialog.msg_parts.Text': ('msg_parts.html#text', 'aidialog/msg_parts.py'), 'aidialog.msg_parts.Text.__init__': ('msg_parts.html#text.__init__', 'aidialog/msg_parts.py'), diff --git a/aidialog/hist.py b/aidialog/hist.py index 350c226..d87285f 100644 --- a/aidialog/hist.py +++ b/aidialog/hist.py @@ -6,12 +6,12 @@ # %% auto #0 __all__ = ['UNSUPPORTED_MSG', 'im_max', 'IMG_TOKS', 'expr_pat', 'expr_mtypes', 'jwrap', 'to_local_time', 'media_item', - 'resize_img', 'output_parts', 'merge_media', 'get_refs', 'sigil_pat', 'get_exprs', 'vars_tag', 'vars_hist', - 'is_nameerr', 'task_tags', 'warning_tag', 'dlg2hist', 'reply2chat', 'dlg2chat', 'chat2dlg', 'reply2dlg', - 'dlg2reply'] + 'resize_img', 'pdf_pages', 'output_parts', 'merge_media', 'get_refs', 'sigil_pat', 'get_exprs', 'vars_tag', + 'vars_hist', 'is_nameerr', 'task_tags', 'warning_tag', 'dlg2hist', 'reply2chat', 'dlg2chat', 'chat2dlg', + 'reply2dlg', 'dlg2reply'] # %% ../nbs/03_hist.ipynb #40ac6466 -import re, ast, base64, binascii, hashlib +import re, ast, base64, binascii, hashlib, zlib from ast import literal_eval from datetime import datetime from zoneinfo import ZoneInfo @@ -114,6 +114,13 @@ def prep_img(self:Message, data, mime, max_im_sz=None): "Resize raster images before they enter LLM context, so they don't consume too many tokens" return resize_img(data, max_im_sz or im_max) if mime in IMG_MIMES else data +# %% ../nbs/03_hist.ipynb #215e15e2 +_pdf_count = re.compile(rb'/Type\s*/Pages[^>]*?/Count\s+(\d+)|/Count\s+(\d+)[^>]*?/Type\s*/Pages', re.S) +def pdf_pages(b:bytes): + "Page count of PDF `b`: the largest `/Count` on a `/Pages` node, in the clear or inside an inflated object stream" + objs = [b] + [zlib.decompressobj().decompress(b[m.end():]) for m in re.finditer(rb'/Type\s*/ObjStm.*?stream\r?\n', b, re.S)] + return max((int(a or c) for o in objs for a,c in _pdf_count.findall(o)), default=0) + # %% ../nbs/03_hist.ipynb #c9792733 def _media_atts(msg, aim_info, max_im_sz=None): "Build media context from a message's attachments that are referenced in the content." diff --git a/aidialog/msg_parts.py b/aidialog/msg_parts.py index 6f8c389..19c858f 100644 --- a/aidialog/msg_parts.py +++ b/aidialog/msg_parts.py @@ -7,10 +7,10 @@ # %% auto #0 __all__ = ['PartType', 'tool_info', 'usage_info', 'think_start', 'think_end', 're_think', 'fence_call_re', 'Part', 'Text', 'Thinking', 'Refusal', 'Media', 'InputImage', 'InputAudio', 'InputVideo', 'InputFile', 'mk_part', 'Msg', - 'ToolUse', 'ToolResult', 'ServerToolResult', 'display_list', 'Completion', 'mk_tool_res_msg', 'sys_text', - 'part_txt', 'data_url', 'url_mime', 'MediaUrl', 'mk_content', 'parse_tools', 'strip_tools', 'conv_tools', - 'extract_fence_call', 'mk_result_fence', 'split_fence_msgs', 'tool_text', 'fmt2hist', 'ToolResponse', - 'StopResponse', 'FullResponse', 'trunc_str', 'mk_tr_details', 'hist2fmt', 'mk_msg', 'mk_msgs'] + 'ToolUse', 'ToolResult', 'display_list', 'Completion', 'mk_tool_res_msg', 'sys_text', 'part_txt', 'data_url', + 'url_mime', 'MediaUrl', 'mk_content', 'parse_tools', 'strip_tools', 'conv_tools', 'extract_fence_call', + 'mk_result_fence', 'split_fence_msgs', 'tool_text', 'fmt2hist', 'ToolResponse', 'StopResponse', + 'FullResponse', 'trunc_str', 'mk_tr_details', 'hist2fmt', 'mk_msg', 'mk_msgs'] # %% ../nbs/00_msg_parts.ipynb #a616b4f5 import base64, json, copy @@ -41,7 +41,7 @@ def replace(self, **kw): return res # %% ../nbs/00_msg_parts.ipynb #5342db55 -PartType = str_enum('PartType', 'text', 'thinking', 'refusal', 'tool_use', 'server_tool_result', 'tool_result', +PartType = str_enum('PartType', 'text', 'thinking', 'refusal', 'tool_use', 'tool_result', 'input_image', 'input_audio', 'input_video', 'input_file') # %% ../nbs/00_msg_parts.ipynb #48496c08 @@ -101,10 +101,11 @@ def _repr_markdown_(self:Part): # %% ../nbs/00_msg_parts.ipynb #ae349c18 class Msg(BasicRepr): - "A normalized message." + "A normalized message; `raw` is the wire form it was parsed from, when a provider produced it." def __init__(self, role, # 'user', 'assistant', or 'tool' - content # list of `Part` + content, # list of `Part` + raw=None # The provider's message as received, replayed verbatim to the same provider ): store_attr() @@ -135,11 +136,6 @@ def __init__(self, id=None, name=None, arguments=None, server=False, text=None, class ToolUse (_ToolPart, tag=PartType.tool_use ): "A tool invocation; `server` marks one the provider ran itself." class ToolResult(_ToolPart, tag=PartType.tool_result): "A tool call's result, `text` holding the output." -class ServerToolResult(Part, tag=PartType.server_tool_result): - "A provider-side tool result, kept as `raw` for round-trips." - def __init__(self, text=None, **kw): - super().__init__(**kw) - store_attr('text') # %% ../nbs/00_msg_parts.ipynb #195b4d89 @patch @@ -370,10 +366,9 @@ def tool_text( # %% ../nbs/00_msg_parts.ipynb #ade55a7b def _extract_tool_parts(d:dict): "Build (tool_use_part, tool_result_part) from a parsed `{.tool}` block" - # Skip server tool calls in deserialization (round trip issues with Gemini/Anthropic) - if not d or d.get('server') or d.get('id') is None: return None - tu = ToolUse (id=d['id'], name=d['name'], arguments=d.get('args') or {}) - tr = ToolResult(id=d['id'], name=d['name'], text=tool_text(d.get('result'))) + if not d or d.get('id') is None: return None + tu = ToolUse (id=d['id'], name=d['name'], arguments=d.get('args') or {}, server=d.get('server', False)) + tr = ToolResult(id=d['id'], name=d['name'], text=tool_text(d.get('result')), server=d.get('server', False)) return tu, tr # %% ../nbs/00_msg_parts.ipynb #916f8df0 @@ -503,7 +498,7 @@ def hist2fmt(msgs:list[Msg], mx=2000, showthink=False)->str: for m in msgs: if m.role == 'assistant': for p in m.content: - if isinstance(p, ToolUse) and not p.server: tus[p.id] = p + if isinstance(p, ToolUse): tus[p.id] = p else: out.append(p.doc(showthink=showthink, mx=mx)) elif m.role == 'tool': for p in m.content: diff --git a/nbs/00_msg_parts.ipynb b/nbs/00_msg_parts.ipynb index 1d648c7..b631fd4 100644 --- a/nbs/00_msg_parts.ipynb +++ b/nbs/00_msg_parts.ipynb @@ -164,7 +164,7 @@ "outputs": [], "source": [ "#| export\n", - "PartType = str_enum('PartType', 'text', 'thinking', 'refusal', 'tool_use', 'server_tool_result', 'tool_result',\n", + "PartType = str_enum('PartType', 'text', 'thinking', 'refusal', 'tool_use', 'tool_result',\n", " 'input_image', 'input_audio', 'input_video', 'input_file')" ] }, @@ -349,10 +349,11 @@ "source": [ "#| export\n", "class Msg(BasicRepr):\n", - " \"A normalized message.\"\n", + " \"A normalized message; `raw` is the wire form it was parsed from, when a provider produced it.\"\n", " def __init__(self,\n", " role, # 'user', 'assistant', or 'tool'\n", - " content # list of `Part`\n", + " content, # list of `Part`\n", + " raw=None # The provider's message as received, replayed verbatim to the same provider\n", " ):\n", " store_attr()\n", "\n", @@ -486,12 +487,7 @@ " store_attr('id,name,arguments,server,text')\n", "\n", "class ToolUse (_ToolPart, tag=PartType.tool_use ): \"A tool invocation; `server` marks one the provider ran itself.\"\n", - "class ToolResult(_ToolPart, tag=PartType.tool_result): \"A tool call's result, `text` holding the output.\"\n", - "class ServerToolResult(Part, tag=PartType.server_tool_result):\n", - " \"A provider-side tool result, kept as `raw` for round-trips.\"\n", - " def __init__(self, text=None, **kw):\n", - " super().__init__(**kw)\n", - " store_attr('text')" + "class ToolResult(_ToolPart, tag=PartType.tool_result): \"A tool call's result, `text` holding the output.\"" ] }, { @@ -1488,10 +1484,9 @@ "#| export\n", "def _extract_tool_parts(d:dict):\n", " \"Build (tool_use_part, tool_result_part) from a parsed `{.tool}` block\"\n", - " # Skip server tool calls in deserialization (round trip issues with Gemini/Anthropic)\n", - " if not d or d.get('server') or d.get('id') is None: return None\n", - " tu = ToolUse (id=d['id'], name=d['name'], arguments=d.get('args') or {})\n", - " tr = ToolResult(id=d['id'], name=d['name'], text=tool_text(d.get('result')))\n", + " if not d or d.get('id') is None: return None\n", + " tu = ToolUse (id=d['id'], name=d['name'], arguments=d.get('args') or {}, server=d.get('server', False))\n", + " tr = ToolResult(id=d['id'], name=d['name'], text=tool_text(d.get('result')), server=d.get('server', False))\n", " return tu, tr" ] }, @@ -1915,7 +1910,7 @@ " for m in msgs:\n", " if m.role == 'assistant':\n", " for p in m.content:\n", - " if isinstance(p, ToolUse) and not p.server: tus[p.id] = p\n", + " if isinstance(p, ToolUse): tus[p.id] = p\n", " else: out.append(p.doc(showthink=showthink, mx=mx))\n", " elif m.role == 'tool':\n", " for p in m.content:\n", @@ -2083,6 +2078,30 @@ "s" ] }, + { + "cell_type": "markdown", + "id": "bf492d17", + "metadata": {}, + "source": [ + "A server call, one the provider ran itself, has no result message of its own, so it renders as a completed block. Re-parsing gives a call and result pair with `server` kept, so a stored conversation replays the call as an ordinary tool call, and nothing tries to run it again:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b2ab0579", + "metadata": {}, + "outputs": [], + "source": [ + "srv = Msg('assistant', [Text('Let me check.'), ToolUse(id='s1', name='web_search', arguments={'query': 'otters'}, server=True)])\n", + "s = hist2fmt([srv])\n", + "h3 = fmt2hist(s)\n", + "test_eq([m.role for m in h3[:2]], ['assistant', 'tool'])\n", + "test_eq((h3[0].content[1].server, h3[1].content[0].server), (True, True))\n", + "test_eq(hist2fmt(h3[:2]), s)\n", + "Markdown(s)" + ] + }, { "cell_type": "markdown", "id": "0d6af3f7", diff --git a/nbs/03_hist.ipynb b/nbs/03_hist.ipynb index b658a15..ece4a97 100644 --- a/nbs/03_hist.ipynb +++ b/nbs/03_hist.ipynb @@ -36,7 +36,7 @@ "outputs": [], "source": [ "#| export\n", - "import re, ast, base64, binascii, hashlib\n", + "import re, ast, base64, binascii, hashlib, zlib\n", "from ast import literal_eval\n", "from datetime import datetime\n", "from zoneinfo import ZoneInfo\n", @@ -554,6 +554,54 @@ " return resize_img(data, max_im_sz or im_max) if mime in IMG_MIMES else data" ] }, + { + "cell_type": "markdown", + "id": "50a851d9", + "metadata": {}, + "source": [ + "A PDF attachment is budgeted per page, and the count comes from the file's page tree rather than a PDF library. The root `/Pages` node holds the total in `/Count`, and no lower node exceeds it, so the largest `/Count` in the file is the page count. Since PDF 1.5 that node may sit inside a compressed object stream, which `zlib` unpacks. An edit that removed pages leaves the old, larger root behind, so the count can only err high, which is the safe side for a token estimate:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "215e15e2", + "metadata": {}, + "outputs": [], + "source": [ + "#| export\n", + "_pdf_count = re.compile(rb'/Type\\s*/Pages[^>]*?/Count\\s+(\\d+)|/Count\\s+(\\d+)[^>]*?/Type\\s*/Pages', re.S)\n", + "def pdf_pages(b:bytes):\n", + " \"Page count of PDF `b`: the largest `/Count` on a `/Pages` node, in the clear or inside an inflated object stream\"\n", + " objs = [b] + [zlib.decompressobj().decompress(b[m.end():]) for m in re.finditer(rb'/Type\\s*/ObjStm.*?stream\\r?\\n', b, re.S)]\n", + " return max((int(a or c) for o in objs for a,c in _pdf_count.findall(o)), default=0)" + ] + }, + { + "cell_type": "markdown", + "id": "d4d63446", + "metadata": {}, + "source": [ + "Two hand-built files show both placements: a page tree in the clear, and the same tree packed into an object stream:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ca3a6979", + "metadata": {}, + "outputs": [], + "source": [ + "def _pdf(objs): return b'%PDF-1.5\\n' + b''.join(f'{i} 0 obj\\n'.encode()+o+b'\\nendobj\\n' for i,o in enumerate(objs, 1)) + b'%%EOF'\n", + "pages = b'<< /Type /Pages /Count 3 /Kids [3 0 R 4 0 R 5 0 R] >>'\n", + "kids = [b'<< /Type /Page /Parent 2 0 R >>']*3\n", + "plain = _pdf([b'<< /Type /Catalog /Pages 2 0 R >>', pages, *kids])\n", + "packed = zlib.compress(b'2 0 ' + pages)\n", + "objstm = _pdf([b'<< /Type /Catalog /Pages 2 0 R >>', b'<< /Type /ObjStm /N 1 /First 4 /Filter /FlateDecode /Length %d >>\\nstream\\n' % len(packed) + packed + b'\\nendstream', *kids])\n", + "test_eq((pdf_pages(plain), pdf_pages(objstm), pdf_pages(b'%PDF-1.5\\n%%EOF')), (3, 3, 0))\n", + "pdf_pages(objstm)" + ] + }, { "cell_type": "code", "execution_count": null,