Merge pull request #311 from SuperMarioYL/fix/response-validation-remaining-approaches

Guard against empty/None/truncated provider responses in rto, self_consistency, reread, leap

Author: codelion

.github/workflows/test.yml

@@ -36,6 +36,14 @@ jobs:
         pip install -r tests/requirements.txt
         pip install -e .
     
+    - name: Cache HuggingFace models
+      uses: actions/cache@v3
+      with:
+        path: ~/.cache/huggingface
+        key: ${{ runner.os }}-hf-codelion-dhara-250m
+        restore-keys: |
+          ${{ runner.os }}-hf-
+
     - name: Run unit tests (no server required)
       run: |
         # Set up local inference environment
@@ -81,10 +89,18 @@ jobs:
         pip install -r tests/requirements.txt
         pip install -e .
     
+    - name: Cache HuggingFace models
+      uses: actions/cache@v3
+      with:
+        path: ~/.cache/huggingface
+        key: ${{ runner.os }}-hf-codelion-dhara-250m
+        restore-keys: |
+          ${{ runner.os }}-hf-
+
     - name: Start optillm server
       run: |
         echo "Starting optillm server for integration tests..."
-        OPTILLM_API_KEY=optillm python optillm.py --model Qwen/Qwen2.5-Coder-0.5B-Instruct --port 8000 &
+        OPTILLM_API_KEY=optillm python optillm.py --model codelion/dhara-250m --port 8000 &
         echo $! > server.pid
         
         # Wait for server to be ready
@@ -174,12 +190,20 @@ jobs:
         pip install -r tests/requirements.txt
         pip install -e .
 
+    - name: Cache HuggingFace models
+      uses: actions/cache@v3
+      with:
+        path: ~/.cache/huggingface
+        key: ${{ runner.os }}-hf-codelion-dhara-250m
+        restore-keys: |
+          ${{ runner.os }}-hf-
+
     - name: Start optillm server with conversation logging
       run: |
         echo "Starting optillm server with conversation logging..."
         mkdir -p /tmp/optillm_conversations
         OPTILLM_API_KEY=optillm python optillm.py \
-          --model Qwen/Qwen2.5-Coder-0.5B-Instruct \
+          --model codelion/dhara-250m \
           --port 8000 \
           --log-conversations \
           --conversation-log-dir /tmp/optillm_conversations &

optillm/__init__.py

@@ -1,5 +1,17 @@
 # Version information
-__version__ = "0.3.15"
+__version__ = "0.3.16"
+
+import os as _os
+
+# An empty Hugging Face token is not the same as "no token". huggingface_hub 1.x
+# (pulled in by transformers 5) sends an empty token as the literal header
+# "Authorization: Bearer ", which the HTTP stack rejects with
+# "Illegal header value b'Bearer '". This happens whenever HF_TOKEN is present but
+# blank, e.g. a forked-PR CI run where ${{ secrets.HF_TOKEN }} resolves to "".
+# Treat any blank HF token env var as unset so anonymous access is used instead.
+for _hf_token_var in ("HF_TOKEN", "HUGGING_FACE_HUB_TOKEN", "HUGGINGFACE_HUB_TOKEN", "HF_HUB_TOKEN"):
+    if _hf_token_var in _os.environ and not _os.environ[_hf_token_var].strip():
+        del _os.environ[_hf_token_var]
 
 # Import from server module
 from .server import (

optillm/leap.py

@@ -28,6 +28,21 @@ def extract_output(self, text: str) -> str:
         match = re.search(r'<output>(.*?)(?:</output>|$)', text, re.DOTALL)
         return match.group(1).strip() if match else ""
 
+    def _extract_content(self, response, context: str) -> str:
+        """Validate a provider response and return its message content.
+
+        Guards against the empty / None / length-truncated responses that would
+        otherwise raise IndexError/TypeError downstream (e.g. when the content
+        is fed to extract_output or split). Mirrors the response-validation
+        idiom already used in moa/bon/plansearch.
+        """
+        if (response is None or
+                not response.choices or
+                response.choices[0].message.content is None or
+                response.choices[0].finish_reason == "length"):
+            raise Exception(f"LEAP: provider returned an empty, None, or truncated response while {context}")
+        return response.choices[0].message.content
+
     def extract_examples_from_query(self, initial_query: str) -> List[Tuple[str, str]]:
         logger.info("Extracting examples from initial query")
 
@@ -66,7 +81,7 @@ def extract_examples_from_query(self, initial_query: str) -> List[Tuple[str, str
             optillm.conversation_logger.log_provider_call(self.request_id, provider_request, response_dict)
 
         self.leap_completion_tokens += response.usage.completion_tokens
-        examples_str = self.extract_output(response.choices[0].message.content)
+        examples_str = self.extract_output(self._extract_content(response, "extracting examples from the query"))
         logger.debug(f"Extracted examples: {examples_str}")
         examples = []
         if examples_str:
@@ -109,7 +124,7 @@ def generate_mistakes(self, examples: List[Tuple[str, str]]) -> List[Tuple[str,
                 response_dict = response.model_dump() if hasattr(response, 'model_dump') else response
                 optillm.conversation_logger.log_provider_call(self.request_id, provider_request, response_dict)
             self.leap_completion_tokens += response.usage.completion_tokens
-            generated_reasoning = response.choices[0].message.content
+            generated_reasoning = self._extract_content(response, "generating reasoning for a mistake example")
             generated_answer = self.extract_output(generated_reasoning)
             if generated_answer != correct_answer:
                 mistakes.append((question, generated_reasoning, generated_answer, correct_answer))
@@ -148,7 +163,7 @@ def generate_low_level_principles(self, mistakes: List[Tuple[str, str, str, str]
                 response_dict = response.model_dump() if hasattr(response, 'model_dump') else response
                 optillm.conversation_logger.log_provider_call(self.request_id, provider_request, response_dict)
             self.leap_completion_tokens += response.usage.completion_tokens
-            self.low_level_principles.append(self.extract_output(response.choices[0].message.content))
+            self.low_level_principles.append(self.extract_output(self._extract_content(response, "generating low-level principles")))
         return self.low_level_principles
 
     def generate_high_level_principles(self) -> List[str]:
@@ -181,7 +196,7 @@ def generate_high_level_principles(self) -> List[str]:
             response_dict = response.model_dump() if hasattr(response, 'model_dump') else response
             optillm.conversation_logger.log_provider_call(self.request_id, provider_request, response_dict)
         self.leap_completion_tokens += response.usage.completion_tokens
-        self.high_level_principles = self.extract_output(response.choices[0].message.content).split("\n")
+        self.high_level_principles = self.extract_output(self._extract_content(response, "generating high-level principles")).split("\n")
         return self.high_level_principles
 
     def apply_principles(self, query: str) -> str:
@@ -210,7 +225,7 @@ def apply_principles(self, query: str) -> str:
             response_dict = response.model_dump() if hasattr(response, 'model_dump') else response
             optillm.conversation_logger.log_provider_call(self.request_id, provider_request, response_dict)
         self.leap_completion_tokens += response.usage.completion_tokens
-        return response.choices[0].message.content
+        return self._extract_content(response, "applying principles to the query")
 
     def solve(self, initial_query: str) -> str:
         logger.info("Starting LEAP process")

optillm/reread.py

@@ -50,11 +50,20 @@ def re2_approach(system_prompt, initial_query, client, model, n=1, request_confi
             response_dict = response.model_dump() if hasattr(response, 'model_dump') else response
             optillm.conversation_logger.log_provider_call(request_id, provider_request, response_dict)
 
+        if response is None or not response.choices:
+            raise Exception("re2_approach: provider returned an empty or None response")
         re2_completion_tokens += response.usage.completion_tokens
         if n == 1:
+            if (response.choices[0].message.content is None or
+                    response.choices[0].finish_reason == "length"):
+                raise Exception("re2_approach: provider returned a None or truncated response")
             return response.choices[0].message.content.strip(), re2_completion_tokens
         else:
-            return [choice.message.content.strip() for choice in response.choices], re2_completion_tokens
+            # Keep only choices that produced usable, non-truncated content,
+            # consistent with how the single-choice path treats a truncation.
+            return [choice.message.content.strip() for choice in response.choices
+                    if choice.message.content is not None
+                    and choice.finish_reason != "length"], re2_completion_tokens
 
     except Exception as e:
         logger.error(f"Error in RE2 approach: {str(e)}")

optillm/rto.py

@@ -41,6 +41,11 @@ def round_trip_optimization(system_prompt: str, initial_query: str, client, mode
         response_dict = response_c1.model_dump() if hasattr(response_c1, 'model_dump') else response_c1
         optillm.conversation_logger.log_provider_call(request_id, provider_request, response_dict)
 
+    if (response_c1 is None or
+            not response_c1.choices or
+            response_c1.choices[0].message.content is None or
+            response_c1.choices[0].finish_reason == "length"):
+        raise Exception("RTO: provider returned an empty, None, or truncated response for the initial code (C1)")
     c1 = response_c1.choices[0].message.content
     rto_completion_tokens += response_c1.usage.completion_tokens
 
@@ -61,6 +66,11 @@ def round_trip_optimization(system_prompt: str, initial_query: str, client, mode
         response_dict = response_q2.model_dump() if hasattr(response_q2, 'model_dump') else response_q2
         optillm.conversation_logger.log_provider_call(request_id, provider_request, response_dict)
 
+    if (response_q2 is None or
+            not response_q2.choices or
+            response_q2.choices[0].message.content is None or
+            response_q2.choices[0].finish_reason == "length"):
+        raise Exception("RTO: provider returned an empty, None, or truncated response for the description (Q2)")
     q2 = response_q2.choices[0].message.content
     rto_completion_tokens += response_q2.usage.completion_tokens
 
@@ -81,6 +91,11 @@ def round_trip_optimization(system_prompt: str, initial_query: str, client, mode
         response_dict = response_c2.model_dump() if hasattr(response_c2, 'model_dump') else response_c2
         optillm.conversation_logger.log_provider_call(request_id, provider_request, response_dict)
 
+    if (response_c2 is None or
+            not response_c2.choices or
+            response_c2.choices[0].message.content is None or
+            response_c2.choices[0].finish_reason == "length"):
+        raise Exception("RTO: provider returned an empty, None, or truncated response for the second code (C2)")
     c2 = response_c2.choices[0].message.content
     rto_completion_tokens += response_c2.usage.completion_tokens
 
@@ -106,6 +121,11 @@ def round_trip_optimization(system_prompt: str, initial_query: str, client, mode
         response_dict = response_c3.model_dump() if hasattr(response_c3, 'model_dump') else response_c3
         optillm.conversation_logger.log_provider_call(request_id, provider_request, response_dict)
 
+    if (response_c3 is None or
+            not response_c3.choices or
+            response_c3.choices[0].message.content is None or
+            response_c3.choices[0].finish_reason == "length"):
+        raise Exception("RTO: provider returned an empty, None, or truncated response for the final code (C3)")
     c3 = response_c3.choices[0].message.content
     rto_completion_tokens += response_c3.usage.completion_tokens
 

optillm/self_consistency.py

@@ -39,6 +39,15 @@ def generate_responses(self, system_prompt: str, user_prompt: str) -> List[str]:
                 response_dict = response.model_dump() if hasattr(response, 'model_dump') else response
                 optillm.conversation_logger.log_provider_call(self.request_id, provider_request, response_dict)
 
+            # Skip empty, None, or length-truncated samples so they do not flow
+            # into similarity clustering (a None content crashes SequenceMatcher).
+            if (response is None or
+                    not response.choices or
+                    response.choices[0].message.content is None or
+                    response.choices[0].finish_reason == "length"):
+                logger.warning("Self-consistency sample was empty, None, or truncated, skipping")
+                continue
+
             self.self_consistency_completion_tokens += response.usage.completion_tokens
             responses.append(response.choices[0].message.content)
         return responses

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "optillm"
-version = "0.3.15"
+version = "0.3.16"
 description = "An optimizing inference proxy for LLMs."
 readme = "README.md"
 license = "Apache-2.0"

requirements.txt

@@ -5,7 +5,7 @@ z3-solver<=4.15.4.0
 aiohttp
 flask
 torch
-transformers
+transformers>=5.0.0
 azure.identity
 tiktoken
 scikit-learn

tests/test.py

@@ -79,7 +79,10 @@ def run_approach(approach_name: str, system_prompt: str, query: str, client, mod
             response = client.chat.completions.create(
                 model=model,
                 messages=messages,
-                temperature=0.7
+                temperature=0.7,
+                # Cap output: the small local test model does not always emit an
+                # EOS token on simple prompts, so bound generation to keep tests fast.
+                max_tokens=64
             )
             result = (response.choices[0].message.content, response.usage.total_tokens)
         else:

tests/test_approaches.py

@@ -16,6 +16,7 @@
 from optillm.cot_reflection import cot_reflection
 from optillm.plansearch import plansearch
 from optillm.leap import leap
+from optillm.rto import round_trip_optimization
 from optillm.mars import multi_agent_reasoning_system
 
 
@@ -34,15 +35,62 @@ class MockChoice:
                     class Message:
                         content = "Test response: 2 + 2 = 4"
                     message = Message()
-                
+                    finish_reason = "stop"
+
                 class MockUsage:
                     completion_tokens = 10
                     total_tokens = 20
-                
+
                 class MockResponse:
                     choices = [MockChoice()]
                     usage = MockUsage()
-                
+
+                return MockResponse()
+
+
+class MockBadClient:
+    """Mock client that simulates degraded provider responses.
+
+    mode:
+      - "empty_choices": response with an empty choices list
+      - "none_content":  a choice whose message.content is None
+      - "length":        a choice truncated with finish_reason == "length"
+    """
+    def __init__(self, mode):
+        self.chat = self.Chat(mode)
+
+    class Chat:
+        def __init__(self, mode):
+            self.completions = self.Completions(mode)
+
+        class Completions:
+            def __init__(self, mode):
+                self.mode = mode
+
+            def create(self, **kwargs):
+                mode = self.mode
+
+                class MockUsage:
+                    completion_tokens = 0
+                    total_tokens = 0
+
+                if mode == "empty_choices":
+                    class MockResponse:
+                        choices = []
+                        usage = MockUsage()
+                    return MockResponse()
+
+                class MockMessage:
+                    content = None if mode == "none_content" else "partial output"
+
+                class MockChoice:
+                    message = MockMessage()
+                    finish_reason = "length" if mode == "length" else "stop"
+
+                class MockResponse:
+                    choices = [MockChoice()]
+                    usage = MockUsage()
+
                 return MockResponse()
 
 
@@ -121,11 +169,50 @@ def test_approach_parameters():
         print(f"✅ {name} has correct parameters")
 
 
+def test_approaches_handle_bad_responses():
+    """Approaches must degrade gracefully on empty / None / length-truncated
+    provider responses instead of raising raw IndexError/TypeError/AttributeError.
+
+    Without the response-validation guards these calls crash with an
+    uncaught IndexError (empty choices) or TypeError/AttributeError (None content
+    flowing into extract_output / SequenceMatcher / .strip()).
+    """
+    system_prompt = "You are a helpful assistant."
+    query = "What is 2 + 2?"
+    model = "mock-model"
+
+    uncontrolled = (IndexError, TypeError, AttributeError)
+    approaches = [
+        ("round_trip_optimization", round_trip_optimization),
+        ("advanced_self_consistency_approach", advanced_self_consistency_approach),
+        ("re2_approach", re2_approach),
+        ("leap", leap),
+    ]
+
+    for mode in ("empty_choices", "none_content", "length"):
+        client = MockBadClient(mode)
+        for name, func in approaches:
+            try:
+                result = func(system_prompt, query, client, model)
+                # Graceful degradation: a normal (content, tokens) tuple.
+                assert isinstance(result, tuple) and len(result) == 2, \
+                    f"{name} ({mode}) returned unexpected {result!r}"
+            except uncontrolled as e:
+                raise AssertionError(
+                    f"{name} ({mode}) raised uncontrolled {type(e).__name__}: {e}")
+            except Exception:
+                # A controlled, informative error is acceptable hardening behaviour.
+                pass
+
+    print("✅ approaches handle bad responses gracefully")
+
+
 if __name__ == "__main__":
     print("Running approach tests...")
-    
+
     test_approach_imports()
     test_basic_approach_calls()
     test_approach_parameters()
-    
+    test_approaches_handle_bad_responses()
+
     print("\nAll tests completed!")