Fix LLM client: handle HF API response format limitations + better error handling"

alpha_factory/infra/llm_client.py

@@ -5,6 +5,7 @@ All outputs are schema-constrained via guided JSON generation.
 """
 import asyncio
 import json
+import re
 from typing import TypeVar
 from pydantic import BaseModel
 from openai import AsyncOpenAI
@@ -19,8 +20,6 @@ class LLMClient:
     - Ollama (local, http://localhost:11434/v1)
     - HuggingFace Inference API (cloud, https://router.huggingface.co/v1)
     - vLLM (local/remote, any OpenAI-compatible endpoint)
-
-    All outputs are JSON-schema-constrained for reliability.
     """
 
     def __init__(self, config: LLMConfig, model_manager=None):
@@ -29,24 +28,19 @@ class LLMClient:
         self._clients: dict[str, AsyncOpenAI] = {}
         self._token_count = 0
 
-    def _get_client(self, base_url: str, api_key: str = "dummy", **headers) -> AsyncOpenAI:
+    def _get_client(self, base_url: str, api_key: str = "dummy", **kwargs) -> AsyncOpenAI:
         """Get or create an AsyncOpenAI client for the given endpoint."""
         key = f"{base_url}|{api_key}"
         if key not in self._clients:
             self._clients[key] = AsyncOpenAI(
                 base_url=base_url,
                 api_key=api_key,
-                default_headers=headers if headers else None,
             )
         return self._clients[key]
 
     def _resolve_model(self, tier: str = "mediumfish", model_override: str | None = None) -> tuple[AsyncOpenAI, str]:
-        """
-        Resolve which client + model to use for a given tier.
-        Priority: model_override > ModelManager selection > config default
-        """
+        """Resolve which client + model to use for a given tier."""
         if model_override:
-            # Direct model name — use default endpoint
             client = self._get_client(self.config.base_url, self.config.api_key)
             return client, model_override
 
@@ -56,7 +50,6 @@ class LLMClient:
             client = self._get_client(base_url, api_key)
             return client, model_name
 
-        # Fallback: use config defaults
         tier_to_model = {
             "microfish": self.config.microfish_model,
             "tinyfish": self.config.tinyfish_model,
@@ -78,26 +71,30 @@ class LLMClient:
     ) -> T:
         """
         Generate a structured JSON response conforming to the given Pydantic schema.
-        Uses guided decoding via response_format.
-
-        Args:
-            prompt: The user prompt
-            schema: Pydantic model class for output validation
-            tier: Model tier (microfish/tinyfish/mediumfish/bigfish)
-            model: Override model name (optional)
-            temperature: Override temperature (optional)
-            system_prompt: System message
+        Tries multiple strategies for JSON output:
+        1. response_format: json_schema (vLLM/Ollama)
+        2. response_format: json_object (some providers)
+        3. Plain text with JSON extraction (fallback for HF)
         """
         client, model_name = self._resolve_model(tier, model)
         temp = temperature or self.config.temperature_generation
         json_schema = schema.model_json_schema()
 
+        # Build the schema instruction to embed in prompt
+        schema_str = json.dumps(json_schema, indent=2)
+        json_instruction = (
+            f"\n\nYou MUST output ONLY valid JSON matching this exact schema. "
+            f"No markdown, no explanation, no ```json blocks — ONLY the raw JSON object.\n"
+            f"Schema:\n{schema_str}"
+        )
+
+        # Strategy 1: Try json_schema format (works with vLLM, newer Ollama)
         try:
             response = await client.chat.completions.create(
                 model=model_name,
                 messages=[
                     {"role": "system", "content": system_prompt},
-                    {"role": "user", "content": prompt},
+                    {"role": "user", "content": prompt + json_instruction},
                 ],
                 temperature=temp,
                 max_tokens=self.config.max_tokens,
@@ -109,32 +106,92 @@ class LLMClient:
                     },
                 },
             )
-        except Exception:
-            # Fallback: some providers don't support json_schema format
-            # Try with json_object format + schema instruction in prompt
-            schema_str = json.dumps(json_schema, indent=2)
-            augmented_prompt = (
-                f"{prompt}\n\n"
-                f"IMPORTANT: Output ONLY valid JSON matching this schema:\n"
-                f"```json\n{schema_str}\n```\n"
-                f"No other text. Just the JSON."
-            )
+            content = response.choices[0].message.content
+            if content and content.strip():
+                self._token_count += response.usage.total_tokens if response.usage else 0
+                return self._parse_json_response(content, schema)
+        except Exception as e:
+            if "json_schema" not in str(e).lower() and "format" not in str(e).lower() and "unsupported" not in str(e).lower():
+                # Not a format issue — try json_object
+                pass
+
+        # Strategy 2: Try json_object format
+        try:
             response = await client.chat.completions.create(
                 model=model_name,
                 messages=[
-                    {"role": "system", "content": system_prompt},
-                    {"role": "user", "content": augmented_prompt},
+                    {"role": "system", "content": system_prompt + "\nAlways respond in valid JSON."},
+                    {"role": "user", "content": prompt + json_instruction},
                 ],
                 temperature=temp,
                 max_tokens=self.config.max_tokens,
                 response_format={"type": "json_object"},
             )
+            content = response.choices[0].message.content
+            if content and content.strip():
+                self._token_count += response.usage.total_tokens if response.usage else 0
+                return self._parse_json_response(content, schema)
+        except Exception:
+            pass
 
+        # Strategy 3: Plain text with JSON extraction (works everywhere including HF)
+        response = await client.chat.completions.create(
+            model=model_name,
+            messages=[
+                {"role": "system", "content": system_prompt + "\nYou always respond with valid JSON only. No other text."},
+                {"role": "user", "content": prompt + json_instruction},
+            ],
+            temperature=temp,
+            max_tokens=self.config.max_tokens,
+        )
         content = response.choices[0].message.content
         self._token_count += response.usage.total_tokens if response.usage else 0
 
-        # Parse and validate
-        data = json.loads(content)
+        if not content or not content.strip():
+            raise ValueError(f"Empty response from model {model_name}")
+
+        return self._parse_json_response(content, schema)
+
+    def _parse_json_response(self, content: str, schema: type[T]) -> T:
+        """
+        Parse JSON from LLM response, handling common issues:
+        - Markdown code blocks (```json ... ```)
+        - Leading/trailing text
+        - Thinking tags (<think>...</think>)
+        """
+        text = content.strip()
+
+        # Remove thinking tags (Qwen/DeepSeek R1 style)
+        text = re.sub(r'<think>.*?</think>', '', text, flags=re.DOTALL).strip()
+
+        # Remove markdown code blocks
+        if "```json" in text:
+            match = re.search(r'```json\s*(.*?)\s*```', text, re.DOTALL)
+            if match:
+                text = match.group(1)
+        elif "```" in text:
+            match = re.search(r'```\s*(.*?)\s*```', text, re.DOTALL)
+            if match:
+                text = match.group(1)
+
+        # Try to find JSON object in the text
+        if not text.startswith('{'):
+            # Look for first { and last }
+            start = text.find('{')
+            end = text.rfind('}')
+            if start != -1 and end != -1 and end > start:
+                text = text[start:end + 1]
+
+        # Parse
+        try:
+            data = json.loads(text)
+        except json.JSONDecodeError as e:
+            raise ValueError(
+                f"Failed to parse JSON from model response.\n"
+                f"Error: {e}\n"
+                f"Response (first 500 chars): {content[:500]}"
+            )
+
         return schema.model_validate(data)
 
     async def generate_text(
@@ -146,7 +203,7 @@ class LLMClient:
         system_prompt: str = "You are a quantitative finance expert.",
         max_tokens: int = 2048,
     ) -> str:
-        """Generate free-text response (for memos/reports only, never for expressions)."""
+        """Generate free-text response (for memos/reports only)."""
         client, model_name = self._resolve_model(tier, model)
         temp = temperature or self.config.temperature_critique
 
@@ -162,7 +219,7 @@ class LLMClient:
 
         content = response.choices[0].message.content
         self._token_count += response.usage.total_tokens if response.usage else 0
-        return content
+        return content or ""
 
     async def health_check(self, tier: str = "mediumfish") -> bool:
         """Check if the model endpoint is reachable."""
@@ -170,10 +227,10 @@ class LLMClient:
             client, model_name = self._resolve_model(tier)
             response = await client.chat.completions.create(
                 model=model_name,
-                messages=[{"role": "user", "content": "Say 'ok'"}],
+                messages=[{"role": "user", "content": "Say ok"}],
                 max_tokens=5,
             )
-            return True
+            return bool(response.choices[0].message.content)
         except Exception:
             return False