Stream thinking chunks via reasoning_content and fix eaten tokens

deepseek/client.py

@@ -59,12 +59,20 @@ def _decode_cid(conversation_id: Optional[str]) -> tuple[Optional[str], Optional
     return (session_id or None), parent
 
 
+class Chunk(str):
+    def __new__(cls, value: str, chunk_type: str):
+        obj = str.__new__(cls, value)
+        obj.chunk_type = chunk_type
+        return obj
+
+
 @dataclass
 class Reply:
     """A completed chat reply plus the id to resume the conversation."""
 
     text: str
     conversation_id: str
+    reasoning_text: Optional[str] = None
 
     def __str__(self) -> str:  # so print(reply) shows the text
         return self.text
@@ -179,8 +187,16 @@ def chat(
         """Return the complete reply (`.text`) plus its `.conversation_id`."""
         s = self.stream(prompt, conversation_id=conversation_id,
                         model=model, thinking=thinking, search=search)
-        text = "".join(s)
-        return Reply(text=text, conversation_id=s.conversation_id)
+        reasoning_chunks = []
+        content_chunks = []
+        for chunk in s:
+            if getattr(chunk, "chunk_type", None) == "THINK":
+                reasoning_chunks.append(chunk)
+            else:
+                content_chunks.append(chunk)
+        text = "".join(content_chunks)
+        reasoning_text = "".join(reasoning_chunks) if reasoning_chunks else None
+        return Reply(text=text, conversation_id=s.conversation_id, reasoning_text=reasoning_text)
 
     def close(self) -> None:
         self._http.close()
@@ -202,7 +218,7 @@ def __init__(self, client: "DeepSeekClient", prompt: str, session_id: str,
         self._search = search
         self._message_id: Optional[int] = None
 
-    def __iter__(self) -> Iterator[str]:
+    def __iter__(self) -> Iterator[Chunk]:
         body = {
             "chat_session_id": self._session_id,
             "parent_message_id": self._parent_id,
@@ -223,7 +239,8 @@ def __iter__(self) -> Iterator[str]:
             "POST", COMPLETION_PATH, json=body, headers=headers
         ) as resp:
             resp.raise_for_status()
-            yield from _parse_sse(resp.iter_lines(), meta)
+            for chunk_type, text in _parse_sse(resp.iter_lines(), meta):
+                yield Chunk(text, chunk_type)
         if meta.get("message_id") is not None:
             self._message_id = meta["message_id"]
 
@@ -232,21 +249,35 @@ def conversation_id(self) -> str:
         return _encode_cid(self._session_id, self._message_id)
 
 
-def _parse_sse(lines, meta: Optional[dict] = None) -> Iterator[str]:
-    """Turn DeepSeek's SSE completion stream into reply-text deltas.
+def _find_fragments(data) -> Iterator[dict]:
+    """Recursively yield all fragment dicts containing 'type' and 'id'."""
+    if isinstance(data, dict):
+        if "type" in data and "id" in data:
+            yield data
+        else:
+            for val in data.values():
+                yield from _find_fragments(val)
+    elif isinstance(data, list):
+        for item in data:
+            yield from _find_fragments(item)
+
+
+def _parse_sse(lines, meta: Optional[dict] = None) -> Iterator[tuple[str, str]]:
+    """Turn DeepSeek's SSE completion stream into (chunk_type, text) tuples.
 
     The stream sends an initial snapshot frame whose `v` is the full response
     object (with `fragments[].content`), then a series of append frames:
       * {"p":"response/fragments/-1/content","o":"APPEND","v":" what"}  (sets path)
       * {"v":"'s"}                                                       (appends to it)
-    We track the active append path and emit only RESPONSE-fragment text.
+    We track the active append path and emit fragment text for all fragments (THINK, RESPONSE, etc.).
 
     If `meta` is given, the assistant's `message_id` is recorded into it (used to
     build the resumable conversation_id). The exact field location can vary, so
     we look in a few plausible spots defensively.
     """
     active_path: Optional[str] = None
-    emitted_initial = False
+    active_type: str = "RESPONSE"
+    processed_fragments = set()
 
     for line in lines:
         if not line or not line.startswith("data:"):
@@ -261,32 +292,33 @@ def _parse_sse(lines, meta: Optional[dict] = None) -> Iterator[str]:
 
         v = obj.get("v")
 
-        # Snapshot frame: full response object.
+        # Snapshot frame: capture message ID.
         if isinstance(v, dict) and "response" in v:
             if meta is not None:
                 _capture_message_id(meta, v)
-            for frag in v["response"].get("fragments", []):
-                if frag.get("type") == "RESPONSE" and frag.get("content"):
-                    active_path = "response/fragments/-1/content"
-                    if not emitted_initial:
-                        emitted_initial = True
-                        yield frag["content"]
-            continue
 
-        # Path-setting append frame.
+        # Update active path if path is specified.
         if "p" in obj:
             active_path = obj["p"]
             if meta is not None and active_path.endswith("message_id") \
                     and isinstance(v, int):
                 meta["message_id"] = v
-            if obj.get("o") == "APPEND" and isinstance(v, str) \
-                    and active_path.endswith("content"):
-                yield v
-            continue
 
-        # Bare append to the current path.
+        # Process any new fragments found in the object.
+        for frag in _find_fragments(obj):
+            frag_id = frag.get("id")
+            frag_type = frag.get("type")
+            content = frag.get("content")
+            if frag_id not in processed_fragments:
+                processed_fragments.add(frag_id)
+                if frag_type:
+                    active_type = frag_type
+                if content:
+                    yield active_type, content
+
+        # Yield any text updates to content.
         if isinstance(v, str) and active_path and active_path.endswith("content"):
-            yield v
+            yield active_type, v
 
 
 def _capture_message_id(meta: dict, snapshot: dict) -> None:

server/api.py

@@ -141,4 +141,4 @@ def gen():
     except Exception as e:
         return _error(f"DeepSeek request failed: {e}")
 
-    return completion_response(req.model, reply.text, prompt, reply.conversation_id)
+    return completion_response(req.model, reply.text, prompt, reply.conversation_id, reasoning_content=reply.reasoning_text)

server/openai_format.py

@@ -62,13 +62,16 @@ def _est_tokens(text: str) -> int:
 
 
 def completion_response(model: str, content: str, prompt: str,
-                        conversation_id: str = None) -> dict:
+                        conversation_id: str = None, reasoning_content: str = None) -> dict:
     """A full (non-streaming) OpenAI chat.completion object.
 
     `conversation_id` is an extra top-level field (outside OpenAI's schema) you
     send back to resume the conversation.
     """
     pt, ct = _est_tokens(prompt), _est_tokens(content)
+    message = {"role": "assistant", "content": content}
+    if reasoning_content:
+        message["reasoning_content"] = reasoning_content
     return {
         "id": _id(),
         "object": "chat.completion",
@@ -78,7 +81,7 @@ def completion_response(model: str, content: str, prompt: str,
         "choices": [
             {
                 "index": 0,
-                "message": {"role": "assistant", "content": content},
+                "message": message,
                 "finish_reason": "stop",
             }
         ],
@@ -114,7 +117,11 @@ def frame(delta: dict, finish=None, extra: dict = None) -> str:
     yield frame({"role": "assistant", "content": ""})
     for d in stream:
         if d:
-            yield frame({"content": d})
+            chunk_type = getattr(d, "chunk_type", "RESPONSE")
+            if chunk_type == "THINK":
+                yield frame({"reasoning_content": d})
+            else:
+                yield frame({"content": d})
     conversation_id = getattr(stream, "conversation_id", None)
     yield frame({}, finish="stop", extra={"conversation_id": conversation_id})
     yield "data: [DONE]\n\n"