fix: remove 'think' content from llm outputs

Certain LLMs (e.g. DeepSeek R1 via Ollama) include 'think' tags and content in outputs, which breaks output JSON parsing and validation. Implemented a fix to remove 'think' tags and content from LLM outputs. Added relevant tests.

Author: SergiiShcherbak

contextgem/internal/__init__.py

@@ -89,6 +89,7 @@
     _is_json_serializable,
     _llm_call_result_is_valid,
     _parse_llm_output_as_json,
+    _remove_thinking_content_from_llm_output,
     _run_async_calls,
     _run_sync,
     _setup_jinja2_template,
@@ -164,6 +165,7 @@
     "_is_json_serializable",
     "_get_sat_model",
     "_setup_jinja2_template",
+    "_remove_thinking_content_from_llm_output",
     # Converters
     # DOCX
     "WORD_XML_NAMESPACES",

contextgem/internal/base/llms.py

@@ -68,6 +68,7 @@
     _group_instances_by_fields,
     _llm_call_result_is_valid,
     _parse_llm_output_as_json,
+    _remove_thinking_content_from_llm_output,
     _run_async_calls,
     _run_sync,
     _validate_parsed_llm_output,
@@ -1311,7 +1312,9 @@ def merge_usage_data(existing: _LLMUsage | None, new: _LLMUsage) -> _LLMUsage:
             )
             all_usage_data = merge_usage_data(all_usage_data, usage_data)
             extracted_data = _validate_parsed_llm_output(
-                _parse_llm_output_as_json(extracted_data),
+                _parse_llm_output_as_json(
+                    _remove_thinking_content_from_llm_output(extracted_data)
+                ),
                 extracted_item_type=extracted_item_type,
                 justification_provided=add_justifications,
                 references_provided=add_references,

contextgem/internal/utils.py

@@ -426,7 +426,49 @@ def _llm_call_result_is_valid(res: tuple[Any, _LLMUsage] | None) -> bool:
     return True
 
 
-def _parse_llm_output_as_json(output_str: str | dict | list) -> dict | list | None:
+def _remove_thinking_content_from_llm_output(output_str: str | None) -> str | None:
+    """
+    Removes thinking content enclosed in <think></think> tags from the beginning of LLM outputs.
+
+    When using local reasoning LLMs (e.g. DeepSeek R1 in Ollama), the output may include
+    thinking steps enclosed in <think></think> tags at the beginning. This function removes those tags
+    and their content only if they appear at the start of the string, then strips any remaining whitespace.
+
+    This preserves any <think></think> tags that might appear later in the content as part of the
+    actual response.
+
+    :param output_str: The output string from an LLM that may contain thinking content, can be None
+                      if LLM outputs invalid content
+    :type output_str: str | None
+
+    :return: The cleaned string without initial thinking content and extra whitespace, or None if
+             the input was None or an error occurred during processing
+    :rtype: str | None
+    """
+    if output_str is None:
+        return None
+
+    try:
+        # Check if the string starts with <think> tag
+        if output_str.strip().startswith("<think>"):
+            # Find the first closing </think> tag
+            end_tag_pos = output_str.find("</think>")
+            if end_tag_pos != -1:
+                # Remove everything from start to the end of </think> tag
+                cleaned_str = output_str[end_tag_pos + len("</think>") :]
+                # Strip any remaining whitespace
+                cleaned_str = cleaned_str.strip()
+                assert len(cleaned_str) > 0, "Cleaned string is empty"
+                return cleaned_str
+
+        return output_str.strip()
+    except (AssertionError, AttributeError):
+        return None
+
+
+def _parse_llm_output_as_json(
+    output_str: str | dict | list | None,
+) -> dict | list | None:
     """
     Parses the provided LLM-generated output into a JSON-compatible Python object.
 
@@ -436,8 +478,9 @@ def _parse_llm_output_as_json(output_str: str | dict | list) -> dict | list | No
     ` ```json` code block, removing them before parsing.
 
     :param output_str: The output string to parse. It may already be a JSON-parsed Python object,
-        a JSON string, or a string containing a JSON code block marked with ` ```json`.
-    :type output_str: str | dict | list
+        a JSON string, or a string containing a JSON code block marked with ` ```json`. Can be None
+        if LLM outputs invalid content.
+    :type output_str: str | dict | list | None
 
     :return: A dictionary, a list, or `None` if parsing fails or the `output_str` type is invalid.
     :rtype: dict | list | None
@@ -448,8 +491,19 @@ def _parse_llm_output_as_json(output_str: str | dict | list) -> dict | list | No
 
     except json.JSONDecodeError:
         try:
-            # Strip of surrounding ```json, if any
-            answer = output_str.lstrip(r"```json").rstrip(r"```")
+            # Handle markdown code blocks using regex
+
+            answer = output_str.strip()
+
+            # Pattern to match content between ```json
+            # (at string start) and ``` (at string end) markers
+            json_block_pattern = r"^```json\s*([\s\S]*?)\s*```$"
+            match = re.match(json_block_pattern, answer)
+
+            if match:
+                # Get the content between the markers
+                answer = match.group(1).strip()
+
             return json.loads(answer)
         except json.JSONDecodeError:
             return None

contextgem/public/llms.py

@@ -946,7 +946,9 @@ async def _query_llm(
                     chat_completion = await task
             else:
                 chat_completion = await task
-            answer = chat_completion.choices[0].message.content.strip()
+            answer = chat_completion.choices[
+                0
+            ].message.content  # str, or None if invalid response
             usage.input = chat_completion.usage.prompt_tokens
             usage.output = chat_completion.usage.completion_tokens
             llm_call_obj._record_response_timestamp()

dev/requirements/requirements.dev.txt

@@ -7,7 +7,7 @@ aiosignal==1.3.2 ; python_version >= "3.10" and python_version < "3.14"
 alabaster==0.7.16 ; python_version >= "3.10" and python_version < "3.14"
 annotated-types==0.7.0 ; python_version >= "3.10" and python_version < "3.14"
 anyio==4.9.0 ; python_version >= "3.10" and python_version < "3.14"
-argcomplete==3.5.3 ; python_version >= "3.10" and python_version < "3.14"
+argcomplete==3.6.2 ; python_version >= "3.10" and python_version < "3.14"
 async-timeout==5.0.1 ; python_version == "3.10"
 attrs==25.3.0 ; python_version >= "3.10" and python_version < "3.14"
 babel==2.17.0 ; python_version >= "3.10" and python_version < "3.14"
@@ -20,7 +20,7 @@ cfgv==3.4.0 ; python_version >= "3.10" and python_version < "3.14"
 charset-normalizer==3.4.2 ; python_version >= "3.10" and python_version < "3.14"
 click==8.1.8 ; python_version >= "3.10" and python_version < "3.14"
 colorama==0.4.6 ; python_version >= "3.10" and python_version < "3.14"
-commitizen==4.6.0 ; python_version >= "3.10" and python_version < "3.14"
+commitizen==4.6.3 ; python_version >= "3.10" and python_version < "3.14"
 coverage==7.8.0 ; python_version >= "3.10" and python_version < "3.14"
 decli==0.6.2 ; python_version >= "3.10" and python_version < "3.14"
 distlib==0.3.9 ; python_version >= "3.10" and python_version < "3.14"
@@ -33,9 +33,10 @@ filelock==3.18.0 ; python_version >= "3.10" and python_version < "3.14"
 frozenlist==1.6.0 ; python_version >= "3.10" and python_version < "3.14"
 fsspec==2025.3.2 ; python_version >= "3.10" and python_version < "3.14"
 h11==0.16.0 ; python_version >= "3.10" and python_version < "3.14"
+hf-xet==1.1.0 ; python_version >= "3.10" and python_version < "3.14" and (platform_machine == "x86_64" or platform_machine == "amd64" or platform_machine == "arm64" or platform_machine == "aarch64")
 httpcore==1.0.9 ; python_version >= "3.10" and python_version < "3.14"
 httpx==0.28.1 ; python_version >= "3.10" and python_version < "3.14"
-huggingface-hub==0.30.2 ; python_version >= "3.10" and python_version < "3.14"
+huggingface-hub==0.31.1 ; python_version >= "3.10" and python_version < "3.14"
 identify==2.6.10 ; python_version >= "3.10" and python_version < "3.14"
 idna==3.10 ; python_version >= "3.10" and python_version < "3.14"
 imagesize==1.4.1 ; python_version >= "3.10" and python_version < "3.14"
@@ -48,7 +49,7 @@ joblib==1.5.0 ; python_version >= "3.10" and python_version < "3.14"
 jsonschema-specifications==2025.4.1 ; python_version >= "3.10" and python_version < "3.14"
 jsonschema==4.23.0 ; python_version >= "3.10" and python_version < "3.14"
 jupyter-core==5.7.2 ; python_version >= "3.10" and python_version < "3.14"
-litellm==1.68.0 ; python_version >= "3.10" and python_version < "3.14"
+litellm==1.68.1 ; python_version >= "3.10" and python_version < "3.14"
 loguru==0.7.3 ; python_version >= "3.10" and python_version < "3.14"
 markupsafe==3.0.2 ; python_version >= "3.10" and python_version < "3.14"
 mosestokenizer==1.2.1 ; python_version >= "3.10" and python_version < "3.14"
@@ -80,7 +81,7 @@ pandas==2.2.3 ; python_version >= "3.10" and python_version < "3.14"
 pathspec==0.12.1 ; python_version >= "3.10" and python_version < "3.14"
 pip-tools==7.4.1 ; python_version >= "3.10" and python_version < "3.14"
 pip==25.1.1 ; python_version >= "3.10" and python_version < "3.14"
-platformdirs==4.3.7 ; python_version >= "3.10" and python_version < "3.14"
+platformdirs==4.3.8 ; python_version >= "3.10" and python_version < "3.14"
 pluggy==1.5.0 ; python_version >= "3.10" and python_version < "3.14"
 pre-commit==4.2.0 ; python_version >= "3.10" and python_version < "3.14"
 prompt-toolkit==3.0.51 ; python_version >= "3.10" and python_version < "3.14"
@@ -106,12 +107,12 @@ requests==2.32.3 ; python_version >= "3.10" and python_version < "3.14"
 rpds-py==0.24.0 ; python_version >= "3.10" and python_version < "3.14"
 safetensors==0.5.3 ; python_version >= "3.10" and python_version < "3.14"
 scikit-learn==1.6.1 ; python_version >= "3.10" and python_version < "3.14"
-scipy==1.15.2 ; python_version >= "3.10" and python_version < "3.14"
-setuptools==80.2.0 ; python_version >= "3.10" and python_version < "3.14"
+scipy==1.15.3 ; python_version >= "3.10" and python_version < "3.14"
+setuptools==80.3.1 ; python_version >= "3.10" and python_version < "3.14"
 six==1.17.0 ; python_version >= "3.10" and python_version < "3.14"
 skops==0.11.0 ; python_version >= "3.10" and python_version < "3.14"
 sniffio==1.3.1 ; python_version >= "3.10" and python_version < "3.14"
-snowballstemmer==2.2.0 ; python_version >= "3.10" and python_version < "3.14"
+snowballstemmer==3.0.0.1 ; python_version >= "3.10" and python_version < "3.14"
 soupsieve==2.7 ; python_version >= "3.10" and python_version < "3.14"
 sphinx-autodoc-typehints==2.3.0 ; python_version >= "3.10" and python_version < "3.14"
 sphinx-book-theme==1.1.4 ; python_version >= "3.10" and python_version < "3.14"
@@ -145,7 +146,7 @@ tzdata==2025.2 ; python_version >= "3.10" and python_version < "3.14"
 uctools==1.3.0 ; python_version >= "3.10" and python_version < "3.14"
 urllib3==2.4.0 ; python_version >= "3.10" and python_version < "3.14"
 vcrpy @ git+https://github.com/shcherbak-ai/vcrpy.git@cda84a815570a3a7bdf93c413901d61230581831 ; python_version >= "3.10" and python_version < "3.14"
-virtualenv==20.30.0 ; python_version >= "3.10" and python_version < "3.14"
+virtualenv==20.31.2 ; python_version >= "3.10" and python_version < "3.14"
 wcwidth==0.2.13 ; python_version >= "3.10" and python_version < "3.14"
 wheel==0.45.1 ; python_version >= "3.10" and python_version < "3.14"
 win32-setctime==1.2.0 ; python_version >= "3.10" and python_version < "3.14" and sys_platform == "win32"

dev/requirements/requirements.main.txt

@@ -19,17 +19,18 @@ filelock==3.18.0 ; python_version >= "3.10" and python_version < "3.14"
 frozenlist==1.6.0 ; python_version >= "3.10" and python_version < "3.14"
 fsspec==2025.3.2 ; python_version >= "3.10" and python_version < "3.14"
 h11==0.16.0 ; python_version >= "3.10" and python_version < "3.14"
+hf-xet==1.1.0 ; python_version >= "3.10" and python_version < "3.14" and (platform_machine == "x86_64" or platform_machine == "amd64" or platform_machine == "arm64" or platform_machine == "aarch64")
 httpcore==1.0.9 ; python_version >= "3.10" and python_version < "3.14"
 httpx==0.28.1 ; python_version >= "3.10" and python_version < "3.14"
-huggingface-hub==0.30.2 ; python_version >= "3.10" and python_version < "3.14"
+huggingface-hub==0.31.1 ; python_version >= "3.10" and python_version < "3.14"
 idna==3.10 ; python_version >= "3.10" and python_version < "3.14"
 importlib-metadata==8.7.0 ; python_version >= "3.10" and python_version < "3.14"
 jinja2==3.1.6 ; python_version >= "3.10" and python_version < "3.14"
 jiter==0.9.0 ; python_version >= "3.10" and python_version < "3.14"
 joblib==1.5.0 ; python_version >= "3.10" and python_version < "3.14"
 jsonschema-specifications==2025.4.1 ; python_version >= "3.10" and python_version < "3.14"
 jsonschema==4.23.0 ; python_version >= "3.10" and python_version < "3.14"
-litellm==1.68.0 ; python_version >= "3.10" and python_version < "3.14"
+litellm==1.68.1 ; python_version >= "3.10" and python_version < "3.14"
 loguru==0.7.3 ; python_version >= "3.10" and python_version < "3.14"
 markupsafe==3.0.2 ; python_version >= "3.10" and python_version < "3.14"
 mosestokenizer==1.2.1 ; python_version >= "3.10" and python_version < "3.14"
@@ -69,8 +70,8 @@ requests==2.32.3 ; python_version >= "3.10" and python_version < "3.14"
 rpds-py==0.24.0 ; python_version >= "3.10" and python_version < "3.14"
 safetensors==0.5.3 ; python_version >= "3.10" and python_version < "3.14"
 scikit-learn==1.6.1 ; python_version >= "3.10" and python_version < "3.14"
-scipy==1.15.2 ; python_version >= "3.10" and python_version < "3.14"
-setuptools==80.2.0 ; python_version >= "3.10" and python_version < "3.14" and platform_system == "Linux" and platform_machine == "x86_64" or python_version >= "3.12" and python_version < "3.14"
+scipy==1.15.3 ; python_version >= "3.10" and python_version < "3.14"
+setuptools==80.3.1 ; python_version >= "3.10" and python_version < "3.14" and platform_system == "Linux" and platform_machine == "x86_64" or python_version >= "3.12" and python_version < "3.14"
 six==1.17.0 ; python_version >= "3.10" and python_version < "3.14"
 skops==0.11.0 ; python_version >= "3.10" and python_version < "3.14"
 sniffio==1.3.1 ; python_version >= "3.10" and python_version < "3.14"