@@ -35,35 +35,81 @@ local function split_messages(text)
3535 local chunk_lines = {}
3636 local chunk_is_user = true
3737
38+ local blocks = {}
39+ local block = {}
40+
3841 --- Insert message and reset/toggle chunk state. User text is trimmed.
3942 local function add_message ()
40- local text_ = table.concat (chunk_lines , ' \n ' )
41-
4243 table.insert (messages , {
4344 role = chunk_is_user and ' user' or ' assistant' ,
44- content = chunk_is_user and vim . trim ( text_ ) or text_ ,
45+ content = blocks ,
4546 })
4647
47- chunk_lines = {}
48+ blocks = {}
4849 chunk_is_user = not chunk_is_user
4950 end
5051
52+ local function add_block ()
53+ if block .type == nil then
54+ block .type = ' text'
55+ end
56+
57+ local text_ = table.concat (chunk_lines , ' \n ' )
58+ -- TODO(ibash) really only want to trim after the <thinking> block, but
59+ -- maybe it's only to trim everything?
60+ -- text_ = chunk_is_user and vim.trim(text_) or text_
61+ text_ = vim .trim (text_ )
62+
63+ if block .type == ' text' then
64+ block .text = text_
65+ elseif block .type == ' thinking' then
66+ block .thinking = text_
67+ elseif block .type == ' redacted_thinking' then
68+ block .data = text_
69+ end
70+
71+ table.insert (blocks , block )
72+
73+ block = {}
74+ chunk_lines = {}
75+ end
76+
5177 for i , line in ipairs (text ) do
52- if i == 1 then
78+ local is_system = i == 1 and line :match (' ^> (.+)' ) ~= nil
79+ -- TODO(ibash) support redacted thinking
80+ local is_thinking = not chunk_is_user and line :match (' <thinking' ) ~= nil
81+ local is_end_thinking = not chunk_is_user
82+ and line :match (' ^</thinking' ) ~= nil
83+ local is_end_message = line == ' ======'
84+
85+ -- print(vim.inspect({
86+ -- i = i,
87+ -- line = line,
88+ -- is_system = is_system,
89+ -- is_thinking = is_thinking,
90+ -- is_end_thinking = is_end_thinking,
91+ -- is_end_message = is_end_message,
92+ -- }))
93+
94+ if is_system then
5395 system = line :match (' ^> (.+)' )
54-
55- if system == nil then
56- table.insert (chunk_lines , line )
57- end
58- elseif line == ' ======' then
96+ elseif is_end_message then
97+ add_block ()
5998 add_message ()
99+ elseif is_thinking then
100+ block .type = ' thinking'
101+ elseif is_end_thinking then
102+ local signature = line :match (' </thinking%s+signature="([^"]*)"' )
103+ block .signature = signature
104+ add_block ()
60105 else
61106 table.insert (chunk_lines , line )
62107 end
63108 end
64109
65110 -- add text after last `======` if not empty
66111 if table.concat (chunk_lines , ' ' ) ~= ' ' then
112+ add_block ()
67113 add_message ()
68114 end
69115
@@ -137,6 +183,30 @@ local function parse_config(text)
137183 end
138184end
139185
186+ local function message_content_to_string (content )
187+ if type (content ) == ' string' then
188+ return content
189+ end
190+
191+ -- TODO(ibash) redacted thinking
192+ local result = {}
193+ for i , block in ipairs (content ) do
194+ if block .type == ' text' then
195+ result [i ] = block .text
196+ elseif block .type == ' thinking' then
197+ result [i ] = ' <thinking>\n '
198+ .. block .thinking
199+ .. ' \n </thinking signature="'
200+ .. block .signature
201+ .. ' ">'
202+ end
203+ end
204+
205+ result = table.concat (result , ' \n\n ' )
206+
207+ return result
208+ end
209+
140210--- Parse a chat file. Must start with a chat name, can follow with a lua table
141211--- of config between `---`. If the next line starts with `> `, it is parsed as
142212--- the system instruction. The rest of the text is parsed as alternating
@@ -184,9 +254,12 @@ function M.to_string(contents, name)
184254 end
185255
186256 if message .role == ' user' then
187- result = result .. ' \n ' .. message .content .. ' \n '
257+ result = result
258+ .. ' \n '
259+ .. message_content_to_string (message .content )
260+ .. ' \n '
188261 else
189- result = result .. message .content
262+ result = result .. message_content_to_string ( message .content )
190263 end
191264 end
192265
0 commit comments