Skip to content

Carry a provider's raw message on Msg, keep server tool calls as ordinary pairs, and count PDF pages without pypdf #48

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
jph00 merged 1 commit into from
Sep 2, 2026
Merged

Carry a provider's raw message on Msg, keep server tool calls as ordinary pairs, and count PDF pages without pypdf #48

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 1 addition & 3 deletions aidialog/_modidx.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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'),
Expand Down Expand Up @@ -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'),
Expand Down
15 changes: 11 additions & 4 deletions aidialog/hist.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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."
Expand Down
29 changes: 12 additions & 17 deletions aidialog/msg_parts.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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()

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand Down
47 changes: 33 additions & 14 deletions nbs/00_msg_parts.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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')"
]
},
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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.\""
]
},
{
Expand Down Expand Up @@ -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"
]
},
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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",
Expand Down
50 changes: 49 additions & 1 deletion nbs/03_hist.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down Expand Up @@ -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,
Expand Down