diff --git a/backend_api_python/app/routes/indicator.py b/backend_api_python/app/routes/indicator.py index e7f3b3308..0d87adf8f 100644 --- a/backend_api_python/app/routes/indicator.py +++ b/backend_api_python/app/routes/indicator.py @@ -1063,13 +1063,12 @@ def _repair_code_via_llm(bad_code: str, validation: Dict[str, Any]) -> str: ) content = (content or "").strip() - if content.startswith("```python"): - content = content[9:] - elif content.startswith("```"): - content = content[3:] - if content.endswith("```"): - content = content[:-3] - return content.strip() or bad_code + # Reuse the hardened stripper: reasoning models prepend prose + # and wrap code in fences with surrounding text; the naive startswith/ + # endswith strip here used to leak both into validation. + from app.services.indicator_ai_generation import _strip_code_fences + + return _strip_code_fences(content) or bad_code def _generate_final_code() -> tuple[str, Dict[str, Any], Dict[str, Any]]: try: diff --git a/backend_api_python/app/services/ai_code_edits.py b/backend_api_python/app/services/ai_code_edits.py index 4f7294916..fa5bc089d 100644 --- a/backend_api_python/app/services/ai_code_edits.py +++ b/backend_api_python/app/services/ai_code_edits.py @@ -37,7 +37,10 @@ def code_edit_user_instruction() -> str: def _parse_payload(response: Any) -> dict[str, Any]: if isinstance(response, dict): return response - raw = str(response or "").strip() + # Reasoning models prepend ... prose; json.loads fails on it. + from app.services.llm import strip_think_blocks + + raw = strip_think_blocks(str(response or "")) if raw.startswith("```json"): raw = raw[7:] elif raw.startswith("```"): diff --git a/backend_api_python/app/services/indicator_ai_generation.py b/backend_api_python/app/services/indicator_ai_generation.py index 495b7afdf..4b8e2737a 100644 --- a/backend_api_python/app/services/indicator_ai_generation.py +++ b/backend_api_python/app/services/indicator_ai_generation.py @@ -4,6 +4,7 @@ import json import os +import re from typing import Any, Callable, Dict, List, Mapping from app.services.ai_code_edits import ( @@ -13,6 +14,7 @@ code_edit_user_instruction, ) from app.services.ai_copilot_context import fit_messages_to_budget +from app.services.llm import strip_think_blocks def _context_block(context: Mapping[str, Any]) -> str: @@ -46,8 +48,19 @@ def _context_block(context: Mapping[str, Any]) -> str: ) +_FENCED_BLOCK_RE = re.compile(r"```(?:python)?\s*\n(.*?)```", re.DOTALL | re.IGNORECASE) + + def _strip_code_fences(content: Any) -> str: - text = str(content or "").strip() + text = strip_think_blocks(str(content or "")) + + # Models sometimes add prose or explanation fences around the code. When any + # complete fenced block exists, return the last one (the final answer). + if "```" in text: + blocks = [b.strip() for b in _FENCED_BLOCK_RE.findall(text) if b.strip()] + if blocks: + return blocks[-1] + if text.startswith("```python"): text = text[9:] elif text.startswith("```"): diff --git a/backend_api_python/app/services/llm.py b/backend_api_python/app/services/llm.py index 3563b4800..2a81db118 100644 --- a/backend_api_python/app/services/llm.py +++ b/backend_api_python/app/services/llm.py @@ -6,6 +6,7 @@ """ import json import os +import re import requests from typing import Dict, Any, Optional, List from enum import Enum @@ -17,6 +18,25 @@ logger = get_logger(__name__) +_THINK_BLOCK_RE = re.compile(r".*?", re.DOTALL | re.IGNORECASE) + + +def strip_think_blocks(text: str) -> str: + """Remove reasoning-model ... chains of thought. + + Reasoning models (MiniMax-M3, DeepSeek-R1, ...) prepend ... + prose to code and JSON answers. Downstream parsers (ast.parse, json.loads) + fail on it, so every LLM response path must strip it before parsing. + An unterminated (truncated response) drops everything from the tag + onward: no usable payload survives truncation anyway. + """ + text = str(text or "") + text = _THINK_BLOCK_RE.sub("", text) + unterminated = text.lower().find("") + if unterminated != -1: + text = text[:unterminated] + return text.strip() + DEFAULT_MAX_TOKENS = 16_384 RETRYABLE_LLM_STATUS_CODES = {408, 409, 425, 429, 500, 502, 503, 504, 529} @@ -1509,8 +1529,8 @@ def safe_call_llm(self, system_prompt: str, user_prompt: str, default_structure: {"role": "user", "content": user_prompt} ], model=model, provider=provider) - # Strip markdown fences if present - clean_text = response_text.strip() + # Strip reasoning-model think prose and markdown fences if present + clean_text = strip_think_blocks(response_text) if clean_text.startswith("```"): first_newline = clean_text.find("\n") if first_newline != -1: diff --git a/backend_api_python/tests/test_indicator_ai_generation_strip.py b/backend_api_python/tests/test_indicator_ai_generation_strip.py new file mode 100644 index 000000000..e07957f05 --- /dev/null +++ b/backend_api_python/tests/test_indicator_ai_generation_strip.py @@ -0,0 +1,138 @@ +"""Regression tests for _strip_code_fences. + +Reasoning models (MiniMax-M3, DeepSeek-R1, ...) prepend ... +chains of thought, and models often wrap code in markdown fences with +surrounding prose. Both shapes used to survive stripping and fail the +ast.parse safety validation with "Unsafe code rejected: Code syntax error". +""" +import ast + +import pytest + +from app.services.indicator_ai_generation import _strip_code_fences + +CODE = "def sma(close, n):\n return close.rolling(window=n).mean()" + + +def assert_valid_python(result: str) -> None: + ast.parse(result) + + +def test_plain_code_unchanged(): + assert _strip_code_fences(CODE) == CODE + + +def test_simple_fence_stripped(): + assert _strip_code_fences(f"```python\n{CODE}\n```") == CODE + + +def test_bare_fence_stripped(): + assert _strip_code_fences(f"```\n{CODE}\n```") == CODE + + +def test_think_block_removed(): + raw = ( + "The user wants an SMA. I will use pandas rolling.\n\n" + f"{CODE}" + ) + result = _strip_code_fences(raw) + assert_valid_python(result) + assert result == CODE + + +def test_think_block_multiline_removed(): + raw = ( + "\nLine one of reasoning.\nLine two.\n\n\n" + f"{CODE}" + ) + result = _strip_code_fences(raw) + assert_valid_python(result) + assert result == CODE + + +def test_think_then_fenced_code(): + raw = ( + "Reasoning here.\n\n" + f"```python\n{CODE}\n```" + ) + result = _strip_code_fences(raw) + assert_valid_python(result) + assert result == CODE + + +def test_unterminated_think_block_truncated(): + raw = f"{CODE}\n\nmodel output was cut off mid-reasoning" + result = _strip_code_fences(raw) + assert_valid_python(result) + assert result == CODE + + +def test_prose_around_fenced_code_returns_last_block(): + raw = ( + "Here is your indicator:\n\n" + f"```python\n{CODE}\n```\n\n" + "This computes a simple moving average." + ) + result = _strip_code_fences(raw) + assert_valid_python(result) + assert result == CODE + + +def test_multiple_fences_returns_last_block(): + draft = "def draft():\n pass\n" + raw = f"```python\n{draft}\n```\n\nRevised:\n\n```python\n{CODE}\n```" + result = _strip_code_fences(raw) + assert result == CODE + + +def test_empty_and_none_inputs(): + assert _strip_code_fences(None) == "" + assert _strip_code_fences("") == "" + assert _strip_code_fences(" ") == "" + + +def test_case_insensitive_think_tag(): + raw = f"reasoning\n\n{CODE}" + result = _strip_code_fences(raw) + assert_valid_python(result) + assert result == CODE + + +def test_prose_only_without_fences_is_returned_as_is(): + # No fenced block and no think tag: behavior is unchanged (the caller's + # validation will reject it), we must not silently invent code. + raw = "Sorry, I cannot generate that indicator." + assert _strip_code_fences(raw) == raw + + +# --------------------------------------------------------------------------- +# Shared strip_think_blocks helper (app.services.llm) +# --------------------------------------------------------------------------- + +from app.services.llm import strip_think_blocks + + +def test_strip_think_blocks_shared_helper(): + assert strip_think_blocks("x\n\ncode()") == "code()" + assert strip_think_blocks("xcode()") == "code()" + assert strip_think_blocks("code()\ntruncated") == "code()" + assert strip_think_blocks("") == "" + assert strip_think_blocks(None) == "" + assert strip_think_blocks("no tags here") == "no tags here" + + +# --------------------------------------------------------------------------- +# Edit-payload JSON parsing must tolerate think prose (app.services.ai_code_edits) +# --------------------------------------------------------------------------- + +from app.services.ai_code_edits import CodeEditError, _parse_payload + + +def test_parse_payload_with_think_prefix(): + raw = 'I will replace the plot line.\n{"operations": []}' + assert _parse_payload(raw) == {"operations": []} + + +def test_parse_payload_still_rejects_garbage(): + with pytest.raises(CodeEditError): + _parse_payload("not json at all")