feat: Enhance Cloudflare integration with fallback handling and improve review planning logic

Author: rostilos

python-ecosystem/inference-orchestrator/src/llm/llm_factory.py

@@ -2,7 +2,7 @@
 import logging
 import json
 from typing import Any, Optional
-from urllib.parse import urlparse
+from urllib.parse import urlparse, urlunparse
 from pydantic import SecretStr
 from langchain_openai import ChatOpenAI
 from langchain_anthropic import ChatAnthropic
@@ -108,6 +108,9 @@ def _normalize_openai_compatible_base_url(ai_base_url: str) -> str:
 
     if _is_cloudflare_base_url(base_url):
         parsed = urlparse(base_url)
+        if parsed.hostname == "api.cloudflare.com" and "/ai/run/" in parsed.path:
+            ai_prefix = parsed.path.split("/ai/run/", 1)[0]
+            return urlunparse(parsed._replace(path=f"{ai_prefix}/ai/v1", params="", query="", fragment=""))
         if parsed.hostname == "api.cloudflare.com" and parsed.path.endswith("/ai"):
             return f"{base_url}/v1"
         return base_url
@@ -155,6 +158,66 @@ def _coerce_openai_compatible_text_content(content: Any) -> str:
     return str(content)
 
 
+_CLOUDFLARE_ROLE_BY_MESSAGE_TYPE = {
+    "human": "user",
+    "ai": "assistant",
+    "system": "system",
+    "tool": "tool",
+    "function": "function",
+}
+
+_CLOUDFLARE_MESSAGE_KEYS = {
+    "role",
+    "content",
+    "name",
+    "tool_calls",
+    "tool_call_id",
+    "function_call",
+}
+
+
+def _cloudflare_message_to_dict(message: Any) -> Any:
+    """Convert dict-like or LangChain message objects into chat message dicts."""
+    if isinstance(message, dict):
+        data = dict(message)
+    else:
+        data = None
+        if hasattr(message, "model_dump"):
+            try:
+                data = message.model_dump(mode="json", exclude_none=True)
+            except TypeError:
+                data = message.model_dump()
+            except Exception:
+                data = None
+        if not isinstance(data, dict) and hasattr(message, "dict"):
+            try:
+                data = message.dict()
+            except Exception:
+                data = None
+        if not isinstance(data, dict):
+            role = getattr(message, "role", None)
+            message_type = getattr(message, "type", None)
+            role = role or _CLOUDFLARE_ROLE_BY_MESSAGE_TYPE.get(str(message_type))
+            content = getattr(message, "content", None)
+            if not role and content is None:
+                return message
+            data = {"role": role, "content": content}
+            for key in ("name", "tool_calls", "tool_call_id", "function_call"):
+                value = getattr(message, key, None)
+                if value:
+                    data[key] = value
+
+    message_type = data.get("type")
+    if not data.get("role") and message_type:
+        data["role"] = _CLOUDFLARE_ROLE_BY_MESSAGE_TYPE.get(str(message_type), str(message_type))
+
+    return {
+        key: value
+        for key, value in data.items()
+        if key in _CLOUDFLARE_MESSAGE_KEYS and value is not None
+    }
+
+
 def _normalize_cloudflare_chat_payload(payload: dict[str, Any]) -> dict[str, Any]:
     """
     Adapt LangChain's OpenAI chat payload to Cloudflare Workers AI's stricter schema.
@@ -167,11 +230,12 @@ def _normalize_cloudflare_chat_payload(payload: dict[str, Any]) -> dict[str, Any
     payload.pop("parallel_tool_calls", None)
 
     messages = payload.get("messages")
-    if not isinstance(messages, list):
+    if not isinstance(messages, (list, tuple)):
         return payload
 
     normalized_messages = []
     for message in messages:
+        message = _cloudflare_message_to_dict(message)
         if not isinstance(message, dict):
             normalized_messages.append(message)
             continue

python-ecosystem/inference-orchestrator/src/service/review/orchestrator/stage_0_planning.py

@@ -6,7 +6,7 @@
 from typing import Any, Dict, Optional
 
 from model.dtos import ReviewRequestDto
-from model.multi_stage import ReviewPlan
+from model.multi_stage import ReviewPlan, FileGroup, ReviewFile
 from utils.prompts.prompt_builder import PromptBuilder
 from utils.diff_processor import ProcessedDiff
 
@@ -75,5 +75,118 @@ async def execute_stage_0_planning(
         content = extract_llm_response_text(response)
         return await parse_llm_response(content, ReviewPlan, llm)
     except Exception as e:
-        logger.error(f"Stage 0 planning failed: {e}")
-        raise ValueError(f"Stage 0 planning failed: {e}")
+        logger.error(f"Stage 0 planning failed, using local fallback plan: {e}")
+        return _build_fallback_review_plan(request, processed_diff)
+
+
+def _build_fallback_review_plan(
+    request: ReviewRequestDto,
+    processed_diff: Optional[ProcessedDiff] = None,
+) -> ReviewPlan:
+    """
+    Build a conservative review plan without another LLM call.
+
+    Stage 0 is an optimization step. If a provider returns empty or malformed
+    planning JSON, the review should still continue with all changed files.
+    """
+    paths = list(dict.fromkeys(request.changedFiles or []))
+    diff_by_path = {df.path: df for df in processed_diff.files} if processed_diff else {}
+
+    if not paths and processed_diff:
+        paths = [df.path for df in processed_diff.files if not df.is_skipped]
+
+    groups: Dict[str, list[ReviewFile]] = {
+        "HIGH": [],
+        "MEDIUM": [],
+        "LOW": [],
+    }
+
+    for path in paths:
+        diff_file = diff_by_path.get(path)
+        priority = _infer_file_priority(path, diff_file)
+        groups[priority].append(
+            ReviewFile(
+                path=path,
+                focus_areas=_infer_focus_areas(path),
+                risk_level=priority,
+                estimated_issues=0,
+            )
+        )
+
+    file_groups = []
+    for priority in ("HIGH", "MEDIUM", "LOW"):
+        files = groups[priority]
+        if not files:
+            continue
+        file_groups.append(
+            FileGroup(
+                group_id=f"FALLBACK_{priority}",
+                priority=priority,
+                rationale="Local fallback plan generated because AI planning output was unavailable",
+                files=files,
+            )
+        )
+
+    return ReviewPlan(
+        analysis_summary=(
+            "Fallback review plan generated locally after AI planning returned "
+            "empty or invalid output."
+        ),
+        file_groups=file_groups,
+        cross_file_concerns=_infer_cross_file_concerns(paths),
+    )
+
+
+def _infer_file_priority(path: str, diff_file: Any = None) -> str:
+    lower = path.lower()
+    if any(marker in lower for marker in (
+        "auth",
+        "security",
+        "permission",
+        "billing",
+        "payment",
+        "migration",
+        "schema",
+        "controller",
+        "handler",
+        "service",
+        "repository",
+    )):
+        return "HIGH"
+    if diff_file and getattr(diff_file, "additions", 0) + getattr(diff_file, "deletions", 0) > 200:
+        return "HIGH"
+    if any(lower.endswith(ext) for ext in (
+        ".md",
+        ".txt",
+        ".json",
+        ".yaml",
+        ".yml",
+        ".toml",
+        ".lock",
+    )):
+        return "LOW"
+    if any(marker in lower for marker in ("/test/", "/tests/", ".test.", ".spec.", "test_")):
+        return "LOW"
+    return "MEDIUM"
+
+
+def _infer_focus_areas(path: str) -> list[str]:
+    lower = path.lower()
+    focus = []
+    if any(marker in lower for marker in ("auth", "security", "permission")):
+        focus.append("SECURITY")
+    if any(marker in lower for marker in ("migration", "schema", "repository", "entity", "model")):
+        focus.append("DATA_ACCESS")
+    if any(marker in lower for marker in ("controller", "handler", "api")):
+        focus.append("API_CONTRACT")
+    if any(marker in lower for marker in ("/test/", "/tests/", ".test.", ".spec.", "test_")):
+        focus.append("TESTING")
+    return focus or ["GENERAL"]
+
+
+def _infer_cross_file_concerns(paths: list[str]) -> list[str]:
+    if len(paths) < 2:
+        return []
+    return [
+        "Check interactions between changed files because AI planning was unavailable."
+    ]

python-ecosystem/inference-orchestrator/tests/test_llm_factory.py

@@ -221,6 +221,14 @@ def test_normalize_cloudflare_ai_gateway_does_not_append_v1(self):
         base = "https://gateway.ai.cloudflare.com/v1/account-id/default/compat"
         assert _normalize_openai_compatible_base_url(base) == base
 
+    def test_normalize_cloudflare_workers_ai_run_endpoint_to_openai_base(self):
+        assert (
+            _normalize_openai_compatible_base_url(
+                "https://api.cloudflare.com/client/v4/accounts/account-id/ai/run/@cf/moonshotai/kimi-k2-instruct"
+            )
+            == "https://api.cloudflare.com/client/v4/accounts/account-id/ai/v1"
+        )
+
     def test_detect_cloudflare_base_url(self):
         assert _is_cloudflare_base_url(
             "https://api.cloudflare.com/client/v4/accounts/id/ai/v1"
@@ -263,6 +271,36 @@ def test_normalize_cloudflare_payload_content_blocks_and_tool_calls(self):
         assert normalized["messages"][3]["content"] == "result"
         assert "parallel_tool_calls" not in normalized
 
+    def test_normalize_cloudflare_payload_langchain_message_objects(self):
+        class MessageObject:
+            type = "human"
+            content = [{"type": "text", "text": "question"}]
+
+        payload = {"messages": (MessageObject(),)}
+
+        normalized = _normalize_cloudflare_chat_payload(payload)
+
+        assert normalized["messages"] == [
+            {"role": "user", "content": "question"}
+        ]
+
+    def test_normalize_cloudflare_payload_model_dump_message(self):
+        class DumpMessage:
+            def model_dump(self, **_kwargs):
+                return {
+                    "type": "system",
+                    "content": [{"type": "text", "text": "sys"}],
+                    "additional_kwargs": {"ignored": True},
+                }
+
+        normalized = _normalize_cloudflare_chat_payload(
+            {"messages": [DumpMessage()]}
+        )
+
+        assert normalized["messages"] == [
+            {"role": "system", "content": "sys"}
+        ]
+
 
 # ── Constants ────────────────────────────────────────────────────
 

python-ecosystem/inference-orchestrator/tests/test_stage_0_branch.py

@@ -54,7 +54,7 @@ async def test_returns_review_plan_from_structured_output(self):
 
     @pytest.mark.asyncio(loop_scope="function")
     async def test_fallback_on_structured_failure(self):
-        """Stage 0 raises ValueError when both structured and raw LLM fail."""
+        """Stage 0 returns a local plan when AI planning fails."""
         mock_llm = MagicMock()
         structured = MagicMock()
         structured.ainvoke = AsyncMock(side_effect=Exception("API error"))
@@ -70,10 +70,41 @@ async def test_fallback_on_structured_failure(self):
         request.enrichmentData = None
         request.projectRules = None
 
-        with pytest.raises(ValueError, match="Stage 0 planning failed"):
-            await execute_stage_0_planning(
-                mock_llm, request, is_incremental=False
-            )
+        result = await execute_stage_0_planning(
+            mock_llm, request, is_incremental=False
+        )
+
+        assert isinstance(result, ReviewPlan)
+        assert result.analysis_summary.startswith("Fallback review plan")
+        assert [f.path for g in result.file_groups for f in g.files] == ["a.py", "b.py"]
+
+    @pytest.mark.asyncio(loop_scope="function")
+    async def test_fallback_on_empty_raw_response(self):
+        """Stage 0 does not fail the review when raw AI output is empty."""
+        mock_llm = MagicMock()
+        structured = MagicMock()
+        structured.ainvoke = AsyncMock(side_effect=Exception("API error"))
+        mock_llm.with_structured_output.return_value = structured
+        raw_response = MagicMock()
+        raw_response.content = ""
+        mock_llm.ainvoke = AsyncMock(return_value=raw_response)
+
+        request = MagicMock()
+        request.changedFiles = ["src/auth/service.py", "README.md"]
+        request.deletedFiles = []
+        request.rawDiff = "diff"
+        request.prTitle = "PR"
+        request.prDescription = "desc"
+        request.enrichmentData = None
+        request.projectRules = None
+
+        result = await execute_stage_0_planning(
+            mock_llm, request, is_incremental=False
+        )
+
+        assert isinstance(result, ReviewPlan)
+        assert result.file_groups[0].priority == "HIGH"
+        assert result.file_groups[-1].priority == "LOW"
 
 
 # ── execute_branch_analysis ──────────────────────────────────────