Show reasoner model's thinking in a <<< thinking section

py/chat.py

@@ -30,7 +30,7 @@ def initialize_chat_window():
         vim.command("redraw")
 
         file_content = vim.eval('trim(join(getline(1, "$"), "\n"))')
-        role_lines = re.findall(r'(^>>> user|^>>> system|^<<< assistant).*', file_content, flags=re.MULTILINE)
+        role_lines = re.findall(r'(^>>> user|^>>> system|^<<< thinking|^<<< assistant).*', file_content, flags=re.MULTILINE)
         if not role_lines[-1].startswith(">>> user"):
             # last role is not user, most likely completion was cancelled before
             vim.command("normal! o")
@@ -58,14 +58,26 @@ def initialize_chat_window():
     try:
         last_content = messages[-1]["content"][-1]
         if last_content['type'] != 'text' or last_content['text']:
-            vim.command("normal! Go\n<<< assistant\n\n")
             vim.command("redraw")
 
             print('Answering...')
-            vim.command("redraw")
+
+            first_thinking_chunk = True
+            first_content_chunk = True
 
             text_chunks = make_chat_text_chunks(messages, options)
-            render_text_chunks(text_chunks)
+            for chunk in text_chunks:
+                if text := chunk.get("thinking"):
+                    if first_thinking_chunk:
+                        first_thinking_chunk = False
+                        vim.command("normal! Go\n<<< thinking\n\n")
+                    vim.command("normal! A" + text)
+                if text := chunk.get("content"):
+                    if first_content_chunk:
+                        first_content_chunk = False
+                        vim.command("normal! Go\n<<< assistant\n\n")
+                    vim.command("normal! A" + text)
+                vim.command("redraw")
 
             vim.command("normal! a\n\n>>> user\n\n")
             vim.command("redraw")

py/complete.py

@@ -34,7 +34,10 @@ def chat_engine(prompt):
         chat_content = f"{initial_prompt}\n\n>>> user\n\n{prompt}".strip()
         messages = parse_chat_messages(chat_content)
         print_debug("[engine-chat] text:\n" + chat_content)
-        return make_chat_text_chunks(messages, config_options)
+        return filter(
+            lambda c: c,
+            map(lambda c: c.get("content"), make_chat_text_chunks(messages, config_options)),
+        )
 
     engines = {"chat": chat_engine, "complete": complete_engine}
 

py/utils.py

@@ -156,10 +156,12 @@ def parse_chat_messages(chat_content):
     for line in lines:
         match line:
             case '>>> system':
-                messages.append({'role': 'system', 'content': [{ 'type': 'text', 'text': '' }]})
+                messages.append({'role': 'system', 'content': ''})
                 current_type = 'system'
+            case '<<< thinking':
+                current_type = 'thinking'
             case '<<< assistant':
-                messages.append({'role': 'assistant', 'content': [{ 'type': 'text', 'text': '' }]})
+                messages.append({'role': 'assistant', 'content': ''})
                 current_type = 'assistant'
             case '>>> user':
                 if messages and messages[-1]['role'] == 'user':
@@ -175,7 +177,9 @@ def parse_chat_messages(chat_content):
                 if not messages:
                     continue
                 match current_type:
-                    case 'assistant' | 'system' | 'user':
+                    case 'system' | 'assistant':
+                        messages[-1]['content'] += '\n' + line
+                    case 'user':
                         messages[-1]['content'][-1]['text'] += '\n' + line
                     case 'include':
                         paths = parse_include_paths(line)
@@ -185,6 +189,10 @@ def parse_chat_messages(chat_content):
 
     for message in messages:
         # strip newlines from the text content as it causes empty responses
+
+        if isinstance(message['content'], str):
+            message['content'] = message['content'].strip()
+            continue
         for content in message['content']:
             if content['type'] == 'text':
                 content['text'] = content['text'].strip()
@@ -334,11 +342,19 @@ def _choices(resp):
 
     def map_chunk_no_stream(resp):
         print_debug("[engine-chat] response: {}", resp)
-        return _choices(resp)[0].get('message', {}).get('content', '')
+        message = _choices(resp)[0].get('message', {})
+        reasoning_content = message.get('reasoning_content', '')
+        content = message.get('content', '')
+        return {"content": content, "thinking": reasoning_content}
 
     def map_chunk_stream(resp):
         print_debug("[engine-chat] response: {}", resp)
-        return _choices(resp)[0].get('delta', {}).get('content', '')
+        delta = _choices(resp)[0].get('delta', {})
+        if reasoning_content := delta.get('reasoning_content'):
+            return {"thinking": reasoning_content}
+        if content := delta.get('content'):
+            return {"content": content}
+        return {"content": ""}
 
     map_chunk = map_chunk_stream if openai_options['stream'] else map_chunk_no_stream
 

syntax/aichat.vim

@@ -1,6 +1,7 @@
 syntax match aichatRole ">>> system"
 syntax match aichatRole ">>> user"
 syntax match aichatRole ">>> include"
+syntax match aichatRole "<<< thinking"
 syntax match aichatRole "<<< assistant"
 
 highlight default link aichatRole Comment