feat(infer): add guided decoding for online and offline models

Introduces a unified utility to parse and validate JSON-structured outputs
from both vLLM and OpenAI API inference results using Pydantic models.
This enables guided decoding for the OnlineLLM.chat method.
The existing guided decoding logic in vllm_infer is refactored to use
this new shared utility, improving consistency and error handling across
inference modes.

weclone/core/inference/offline_infer.py

@@ -6,6 +6,7 @@ from llamafactory.data import get_template_and_fix_tokenizer
 from llamafactory.extras.misc import get_device_count
 from llamafactory.hparams import get_infer_args
 from llamafactory.model import load_tokenizer
+from openai.types.chat import ChatCompletion
 from pydantic import BaseModel
 from vllm import LLM, SamplingParams
 from vllm.lora.request import LoRARequest
@@ -30,6 +31,42 @@ def extract_json_from_text(text: str) -> str:
     return text.strip()
 
 
+def parse_guided_decoding_results(
+    results: List[RequestOutput] | List[ChatCompletion] | List, guided_decoding_class: type[BaseModel]
+) -> tuple[List[Optional[BaseModel]], List[int]]:
+    """Parse guided decoding results and return parsed results with failed indices.
+
+    Args:
+        results: Raw vLLM generation results
+        guided_decoding_class: Pydantic model class for validation
+
+    Returns:
+        tuple: (parsed_results, failed_indices) where failed_indices contains
+               indices of failed JSON parsing
+    """
+    parsed_results = []
+    failed_indexs = []
+
+    for idx, result in enumerate(results):
+        try:
+            if isinstance(result, RequestOutput):
+                json_text = extract_json_from_text(result.outputs[0].text)
+            elif isinstance(result, ChatCompletion):
+                json_text = extract_json_from_text(result.choices[0].message.content)
+            else:
+                raise ValueError(f"Unsupported result type: {type(result)}")
+            parsed_result = guided_decoding_class.model_validate_json(json_text)
+            parsed_results.append(parsed_result)
+        except Exception as e:
+            logger.warning(
+                f"Failed to parse JSON from result at sequence index {idx}: {result.outputs[0].text[:100]}..., error: {e}"
+            )
+            failed_indexs.append(idx)
+            parsed_results.append(None)
+
+    return parsed_results, failed_indexs
+
+
 def vllm_infer(
     inputs: List[str],
     model_name_or_path: str,
@@ -147,22 +184,9 @@ def vllm_infer(
     del llm
     torch.cuda.empty_cache()
 
-    failed_indexs = []
     if guided_decoding_class:
         # TODO better json decode  https://github.com/vllm-project/vllm/commit/1d0ae26c8544fd5a62e171e30c2dcc2973a23bc8#diff-3b27790a2ce97bc50cdd5476f7b0057da682ed0d1ec8426a7b76c5e21454e57d
-        parsed_results = []
-        for idx, result in enumerate(results):
-            try:
-                json_text = extract_json_from_text(result.outputs[0].text)
-                parsed_result = guided_decoding_class.model_validate_json(json_text)
-                parsed_results.append(parsed_result)
-            except Exception as e:
-                # Note that the failed_indexs is the sequential index ID, not the original input ID.
-                logger.warning(
-                    f"Failed to parse JSON from result at sequence index {idx}: {result.outputs[0].text[:100]}..., error: {e}"
-                )
-                failed_indexs.append(idx)
-                parsed_results.append(None)
-        results = parsed_results
-
-    return results, failed_indexs
+        parsed_results, failed_indexs = parse_guided_decoding_results(results, guided_decoding_class)
+        return parsed_results, failed_indexs
+    else:
+        return results, []

weclone/core/inference/online_infer.py

@@ -2,8 +2,11 @@ from concurrent.futures import Future, ThreadPoolExecutor
 from typing import Any, Callable, List, Optional, Union
 
 from openai import OpenAI
-from openai.types.chat import ChatCompletionMessageParam
+from openai.types.chat import ChatCompletion, ChatCompletionMessageParam
+from pydantic import BaseModel
 
+from weclone.core.inference.offline_infer import parse_guided_decoding_results
+from weclone.utils.log import logger
 from weclone.utils.retry import retry_openai_api
 
 
@@ -68,8 +71,23 @@ class OnlineLLM:
         top_p: float = 0.95,
         stream: bool = False,
         callback: Optional[Callable[[int, Any], None]] = None,
-    ) -> List[Union[Any, Exception]]:
-        """Process multiple chat requests concurrently using thread pool"""
+        guided_decoding_class: Optional[type[BaseModel]] = None,
+    ) -> Union[List[Union[ChatCompletion, Exception]], tuple[List[Optional[BaseModel]], List[int]]]:
+        """Process multiple chat requests concurrently using thread pool
+
+        Args:
+            prompts: List of prompt strings
+            temperature: Sampling temperature
+            max_tokens: Maximum tokens to generate
+            top_p: Top-p sampling parameter
+            stream: Whether to stream the response
+            callback: Optional callback function called for each result
+            guided_decoding_class: Pydantic model class for JSON validation
+
+        Returns:
+            If enable_json_decode is False: List of ChatCompletion or Exception objects
+            If enable_json_decode is True: Tuple of (parsed_results, failed_indices)
+        """
         futures = []
 
         for i, prompt in enumerate(prompts):
@@ -89,6 +107,18 @@ class OnlineLLM:
                 if callback:
                     callback(i, e)
 
+        if guided_decoding_class:
+            successful_results = [r for r in results if isinstance(r, ChatCompletion)]
+            if len(successful_results) < len(results):
+                logger.warning(
+                    f"Some requests failed, only {len(successful_results)}/{len(results)} will be parsed"
+                )
+
+            parsed_results, failed_indexs = parse_guided_decoding_results(
+                successful_results, guided_decoding_class
+            )
+            return parsed_results, failed_indexs
+
         return results
 
     def close(self):