diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml
index 792d1bc5..be53dc52 100644
--- a/.github/workflows/publish.yml
+++ b/.github/workflows/publish.yml
@@ -79,48 +79,26 @@ jobs:
type=semver,pattern={{major}}.{{minor}}
type=raw,value=latest
- # Build and push proxy AMD64
- - name: Build and push proxy_only Docker image AMD64
+ # Build and push proxy_only multi-arch
+ - name: Build and push proxy_only Docker image (multi-arch)
uses: docker/build-push-action@v5
with:
context: .
file: Dockerfile.proxy_only
push: true
- platforms: linux/amd64
+ platforms: linux/amd64,linux/arm64
tags: ${{ steps.meta-proxy.outputs.tags }}
labels: ${{ steps.meta-proxy.outputs.labels }}
- cache-from: type=gha,scope=proxy-amd64
- cache-to: type=gha,scope=proxy-amd64,mode=max
+ cache-from: type=gha
+ cache-to: type=gha,mode=max
outputs: type=registry,compression=zstd,compression-level=5
- # Cleanup after AMD64 build
- - name: Cleanup after AMD64 build
+ # Cleanup after proxy build
+ - name: Cleanup after proxy build
run: |
docker system prune -af
docker builder prune -af
df -h
-
- # Build proxy ARM64
- - name: Build and push proxy_only Docker image ARM64
- uses: docker/build-push-action@v5
- with:
- context: .
- file: Dockerfile.proxy_only
- push: true
- platforms: linux/arm64
- tags: ${{ steps.meta-proxy.outputs.tags }}
- labels: ${{ steps.meta-proxy.outputs.labels }}
- cache-from: type=gha,scope=proxy-arm64
- cache-to: type=gha,scope=proxy-arm64,mode=max
- outputs: type=registry,compression=zstd,compression-level=5
-
- # Cleanup after proxy builds
- - name: Cleanup after proxy builds
- run: |
- docker system prune -af
- docker builder prune -af
- find /tmp -type f -user $(id -u) -exec rm -f {} + 2>/dev/null || true
- df -h
# Extract metadata for full image
- name: Extract metadata for Docker
@@ -133,35 +111,15 @@ jobs:
type=semver,pattern={{major}}.{{minor}}
latest
- # Build full image AMD64
- - name: Build and push Docker image AMD64
- uses: docker/build-push-action@v5
- with:
- context: .
- push: true
- platforms: linux/amd64
- tags: ${{ steps.meta.outputs.tags }}
- labels: ${{ steps.meta.outputs.labels }}
- cache-from: type=gha,scope=full-amd64
- cache-to: type=gha,scope=full-amd64,mode=max
- outputs: type=registry,compression=zstd,compression-level=5
-
- # Cleanup between architectures
- - name: Cleanup between architectures
- run: |
- docker system prune -af
- docker builder prune -af
- df -h
-
- # Build full image ARM64
- - name: Build and push Docker image ARM64
+ # Build full image multi-arch
+ - name: Build and push Docker image (multi-arch)
uses: docker/build-push-action@v5
with:
context: .
push: true
- platforms: linux/arm64
+ platforms: linux/amd64,linux/arm64
tags: ${{ steps.meta.outputs.tags }}
labels: ${{ steps.meta.outputs.labels }}
- cache-from: type=gha,scope=full-arm64
- cache-to: type=gha,scope=full-arm64,mode=max
+ cache-from: type=gha
+ cache-to: type=gha,mode=max
outputs: type=registry,compression=zstd,compression-level=5
diff --git a/optillm.py b/optillm.py
index 6ea2d191..e1dc9fab 100644
--- a/optillm.py
+++ b/optillm.py
@@ -93,6 +93,52 @@ def get_config():
default_client = LiteLLMWrapper()
return default_client, API_KEY
+def count_reasoning_tokens(text: str, tokenizer=None) -> int:
+ """
+ Count tokens within ... tags in the given text.
+
+ Args:
+ text: The text to analyze
+ tokenizer: Optional tokenizer instance for precise counting
+
+ Returns:
+ Number of reasoning tokens (0 if no think tags found)
+ """
+ if not text or not isinstance(text, str):
+ return 0
+
+ # Extract all content within ... tags
+ # Handle both complete and truncated think blocks
+
+ # First, find all complete ... blocks
+ complete_pattern = r'(.*?)'
+ complete_matches = re.findall(complete_pattern, text, re.DOTALL)
+
+ # Then check for unclosed tag (truncated response)
+ # This finds that doesn't have a matching after it
+ truncated_pattern = r'(?!.*)(.*)$'
+ truncated_match = re.search(truncated_pattern, text, re.DOTALL)
+
+ # Combine all thinking content
+ thinking_content = ''.join(complete_matches)
+ if truncated_match:
+ thinking_content += truncated_match.group(1)
+
+ if not thinking_content:
+ return 0
+
+ if tokenizer and hasattr(tokenizer, 'encode'):
+ # Use tokenizer for precise counting
+ try:
+ tokens = tokenizer.encode(thinking_content)
+ return len(tokens)
+ except Exception as e:
+ logger.warning(f"Failed to count tokens with tokenizer: {e}")
+
+ # Fallback: rough estimation (4 chars per token on average, minimum 1 token for non-empty content)
+ content_length = len(thinking_content.strip())
+ return max(1, content_length // 4) if content_length > 0 else 0
+
# Server configuration
server_config = {
'approach': 'none',
@@ -678,11 +724,22 @@ def proxy():
if stream:
return Response(generate_streaming_response(response, model), content_type='text/event-stream')
else:
+ # Calculate reasoning tokens from the response
+ reasoning_tokens = 0
+ if isinstance(response, str):
+ reasoning_tokens = count_reasoning_tokens(response)
+ elif isinstance(response, list) and response:
+ # For multiple responses, sum up reasoning tokens from all
+ reasoning_tokens = sum(count_reasoning_tokens(resp) for resp in response if isinstance(resp, str))
+
response_data = {
'model': model,
'choices': [],
'usage': {
'completion_tokens': completion_tokens,
+ 'completion_tokens_details': {
+ 'reasoning_tokens': reasoning_tokens
+ }
}
}
diff --git a/optillm/__init__.py b/optillm/__init__.py
index 4c1aac43..9e0162ea 100644
--- a/optillm/__init__.py
+++ b/optillm/__init__.py
@@ -2,7 +2,7 @@
import os
# Version information
-__version__ = "0.1.22"
+__version__ = "0.1.26"
# Get the path to the root optillm.py
spec = util.spec_from_file_location(
@@ -27,6 +27,7 @@
extract_optillm_approach = module.extract_optillm_approach
get_config = module.get_config
load_plugins = module.load_plugins
+count_reasoning_tokens = module.count_reasoning_tokens
# Export execution functions
execute_single_approach = module.execute_single_approach
@@ -48,6 +49,7 @@
'extract_optillm_approach',
'get_config',
'load_plugins',
+ 'count_reasoning_tokens',
'execute_single_approach',
'execute_combined_approaches',
'execute_parallel_approaches',
diff --git a/optillm/inference.py b/optillm/inference.py
index 94a003db..92b629ee 100644
--- a/optillm/inference.py
+++ b/optillm/inference.py
@@ -18,6 +18,7 @@
import traceback
import platform
import sys
+import re
from optillm.cot_decoding import cot_decode
from optillm.entropy_decoding import entropy_decode
@@ -29,6 +30,52 @@
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
+def count_reasoning_tokens(text: str, tokenizer=None) -> int:
+ """
+ Count tokens within ... tags in the given text.
+
+ Args:
+ text: The text to analyze
+ tokenizer: Optional tokenizer instance for precise counting
+
+ Returns:
+ Number of reasoning tokens (0 if no think tags found)
+ """
+ if not text or not isinstance(text, str):
+ return 0
+
+ # Extract all content within ... tags
+ # Handle both complete and truncated think blocks
+
+ # First, find all complete ... blocks
+ complete_pattern = r'(.*?)'
+ complete_matches = re.findall(complete_pattern, text, re.DOTALL)
+
+ # Then check for unclosed tag (truncated response)
+ # This finds that doesn't have a matching after it
+ truncated_pattern = r'(?!.*)(.*)$'
+ truncated_match = re.search(truncated_pattern, text, re.DOTALL)
+
+ # Combine all thinking content
+ thinking_content = ''.join(complete_matches)
+ if truncated_match:
+ thinking_content += truncated_match.group(1)
+
+ if not thinking_content:
+ return 0
+
+ if tokenizer and hasattr(tokenizer, 'encode'):
+ # Use tokenizer for precise counting
+ try:
+ tokens = tokenizer.encode(thinking_content)
+ return len(tokens)
+ except Exception as e:
+ logger.warning(f"Failed to count tokens with tokenizer: {e}")
+
+ # Fallback: rough estimation (4 chars per token on average, minimum 1 token for non-empty content)
+ content_length = len(thinking_content.strip())
+ return max(1, content_length // 4) if content_length > 0 else 0
+
# MLX Support for Apple Silicon
try:
import mlx.core as mx
@@ -1502,10 +1549,11 @@ def __init__(
self.message.logprobs = logprobs
class ChatCompletionUsage:
- def __init__(self, prompt_tokens: int, completion_tokens: int, total_tokens: int):
+ def __init__(self, prompt_tokens: int, completion_tokens: int, total_tokens: int, reasoning_tokens: int = 0):
self.prompt_tokens = prompt_tokens
self.completion_tokens = completion_tokens
self.total_tokens = total_tokens
+ self.reasoning_tokens = reasoning_tokens
class ChatCompletion:
def __init__(self, response_dict: Dict):
@@ -1547,7 +1595,10 @@ def model_dump(self) -> Dict:
"usage": {
"prompt_tokens": self.usage.prompt_tokens,
"completion_tokens": self.usage.completion_tokens,
- "total_tokens": self.usage.total_tokens
+ "total_tokens": self.usage.total_tokens,
+ "completion_tokens_details": {
+ "reasoning_tokens": getattr(self.usage, 'reasoning_tokens', 0)
+ }
}
}
@@ -1766,7 +1817,7 @@ def create(
logger.debug(f"ThinkDeeper tokens: user={user_max_tokens}, thinking={max_thinking_tokens}, adjusted={adjusted_max_tokens}")
- result = thinkdeeper_decode_mlx(
+ result, reasoning_tokens = thinkdeeper_decode_mlx(
pipeline.model,
pipeline.tokenizer,
messages,
@@ -1774,7 +1825,7 @@ def create(
)
else:
logger.info("Using PyTorch ThinkDeeper implementation")
- result = thinkdeeper_decode(
+ result, reasoning_tokens = thinkdeeper_decode(
pipeline.current_model,
pipeline.tokenizer,
messages,
@@ -1850,6 +1901,11 @@ def create(
prompt_tokens = len(pipeline.tokenizer.encode(prompt))
completion_tokens = sum(token_counts)
+ # Calculate reasoning tokens from all responses
+ total_reasoning_tokens = 0
+ for response in responses:
+ total_reasoning_tokens += count_reasoning_tokens(response, pipeline.tokenizer)
+
# Create OpenAI-compatible response format
response_dict = {
"id": f"chatcmpl-{int(time.time()*1000)}",
@@ -1871,7 +1927,8 @@ def create(
"usage": {
"prompt_tokens": prompt_tokens,
"completion_tokens": completion_tokens,
- "total_tokens": completion_tokens + prompt_tokens
+ "total_tokens": completion_tokens + prompt_tokens,
+ "reasoning_tokens": total_reasoning_tokens
}
}
diff --git a/optillm/plugins/deep_research/research_engine.py b/optillm/plugins/deep_research/research_engine.py
index 2b09f1aa..bfecab8a 100644
--- a/optillm/plugins/deep_research/research_engine.py
+++ b/optillm/plugins/deep_research/research_engine.py
@@ -375,7 +375,7 @@ def decompose_query(self, system_prompt: str, initial_query: str) -> List[str]:
for line in content.split('\n'):
line = line.strip()
if re.match(r'^\d+\.', line):
- query = re.sub(r'^\d+\.\s*', '', line).strip()
+ query = re.sub(r'^\d+\.\s*\[?(.*?)\]?$', r'\1', line).strip()
if query:
queries.append(query)
diff --git a/optillm/plugins/deepthink/__init__.py b/optillm/plugins/deepthink/__init__.py
index cf0c221e..e7922a15 100644
--- a/optillm/plugins/deepthink/__init__.py
+++ b/optillm/plugins/deepthink/__init__.py
@@ -3,4 +3,9 @@
A plugin that combines SELF-DISCOVER framework with uncertainty-routed
chain-of-thought for enhanced reasoning capabilities.
-"""
\ No newline at end of file
+"""
+
+from .self_discover import SelfDiscover
+from .uncertainty_cot import UncertaintyRoutedCoT
+
+__all__ = ['SelfDiscover', 'UncertaintyRoutedCoT']
\ No newline at end of file
diff --git a/optillm/plugins/deepthink_plugin.py b/optillm/plugins/deepthink_plugin.py
index aef2adcf..bdf76021 100644
--- a/optillm/plugins/deepthink_plugin.py
+++ b/optillm/plugins/deepthink_plugin.py
@@ -5,11 +5,9 @@
for enhanced reasoning in large language models.
"""
-import os
-import sys
-import importlib.util
import logging
from typing import Tuple, Dict, Any
+from optillm.plugins.deepthink import SelfDiscover, UncertaintyRoutedCoT
# Plugin identifier for optillm
SLUG = "deepthink"
@@ -41,98 +39,72 @@ def run(
"""
logger.info("Starting Deep Think reasoning process")
- # Get the directory where this plugin is located
- plugin_dir = os.path.dirname(os.path.abspath(__file__))
- deepthink_dir = os.path.join(plugin_dir, 'deepthink')
+ # Extract configuration parameters
+ config = _parse_config(request_config or {})
- # Add the deepthink directory to the Python path temporarily
- if deepthink_dir not in sys.path:
- sys.path.insert(0, deepthink_dir)
+ # Initialize components
+ self_discover = SelfDiscover(
+ client=client,
+ model=model,
+ max_tokens=config["max_tokens"]
+ )
- try:
- # Load the modules dynamically
- self_discover_file = os.path.join(deepthink_dir, 'self_discover.py')
- uncertainty_cot_file = os.path.join(deepthink_dir, 'uncertainty_cot.py')
-
- spec1 = importlib.util.spec_from_file_location("self_discover", self_discover_file)
- self_discover_module = importlib.util.module_from_spec(spec1)
- spec1.loader.exec_module(self_discover_module)
-
- spec2 = importlib.util.spec_from_file_location("uncertainty_cot", uncertainty_cot_file)
- uncertainty_cot_module = importlib.util.module_from_spec(spec2)
- spec2.loader.exec_module(uncertainty_cot_module)
-
- # Extract configuration parameters
- config = _parse_config(request_config or {})
-
- # Initialize components
- self_discover = self_discover_module.SelfDiscover(
- client=client,
- model=model,
- max_tokens=config["max_tokens"]
- )
-
- uncertainty_cot = uncertainty_cot_module.UncertaintyRoutedCoT(
- client=client,
- model=model,
- max_tokens=config["max_tokens"]
- )
-
- total_tokens = 0
-
- # Stage 1: SELF-DISCOVER reasoning structure (if enabled)
- reasoning_structure = None
- if config["enable_self_discover"]:
- logger.info("Discovering task-specific reasoning structure")
-
- discovery_result = self_discover.discover_reasoning_structure(
- task_description=_extract_task_description(initial_query, system_prompt),
- task_examples=None # Could be enhanced to extract examples
- )
-
- reasoning_structure = discovery_result["reasoning_structure"]
- total_tokens += discovery_result["completion_tokens"]
-
- logger.info(f"Discovered reasoning structure with {len(reasoning_structure)} components")
-
- # Prepare enhanced prompt
- enhanced_prompt = _create_enhanced_prompt(
- system_prompt=system_prompt,
- initial_query=initial_query,
- reasoning_structure=reasoning_structure,
- config=config
- )
-
- # Stage 2: Uncertainty-routed generation
- logger.info("Generating response with uncertainty routing")
+ uncertainty_cot = UncertaintyRoutedCoT(
+ client=client,
+ model=model,
+ max_tokens=config["max_tokens"]
+ )
+
+ total_tokens = 0
+
+ # Stage 1: SELF-DISCOVER reasoning structure (if enabled)
+ reasoning_structure = None
+ if config["enable_self_discover"]:
+ logger.info("Discovering task-specific reasoning structure")
- generation_result = uncertainty_cot.generate_with_uncertainty_routing(
- prompt=enhanced_prompt,
- num_samples=config["deepthink_samples"],
- confidence_threshold=config["confidence_threshold"],
- temperature=config["temperature"],
- top_p=config["top_p"]
+ discovery_result = self_discover.discover_reasoning_structure(
+ task_description=_extract_task_description(initial_query, system_prompt),
+ task_examples=None # Could be enhanced to extract examples
)
- total_tokens += generation_result["completion_tokens"]
+ reasoning_structure = discovery_result["reasoning_structure"]
+ total_tokens += discovery_result["completion_tokens"]
- # Log routing decision
- logger.info(f"Routing decision: {generation_result['routing_decision']} "
- f"(confidence: {generation_result['confidence_score']:.3f})")
-
- final_response = generation_result["final_response"]
-
- # Clean up the response if needed
- final_response = _clean_response(final_response)
-
- logger.info(f"Deep Think completed successfully. Total tokens: {total_tokens}")
-
- return final_response, total_tokens
-
- finally:
- # Remove from path after use
- if deepthink_dir in sys.path:
- sys.path.remove(deepthink_dir)
+ logger.info(f"Discovered reasoning structure with {len(reasoning_structure)} components")
+
+ # Prepare enhanced prompt
+ enhanced_prompt = _create_enhanced_prompt(
+ system_prompt=system_prompt,
+ initial_query=initial_query,
+ reasoning_structure=reasoning_structure,
+ config=config
+ )
+
+ # Stage 2: Uncertainty-routed generation
+ logger.info("Generating response with uncertainty routing")
+
+ generation_result = uncertainty_cot.generate_with_uncertainty_routing(
+ prompt=enhanced_prompt,
+ num_samples=config["deepthink_samples"],
+ confidence_threshold=config["confidence_threshold"],
+ temperature=config["temperature"],
+ top_p=config["top_p"]
+ )
+
+ total_tokens += generation_result["completion_tokens"]
+
+ # Log routing decision
+ logger.info(f"Routing decision: {generation_result['routing_decision']} "
+ f"(confidence: {generation_result['confidence_score']:.3f})")
+
+ final_response = generation_result["final_response"]
+
+ # Clean up the response if needed
+ final_response = _clean_response(final_response)
+
+ logger.info(f"Deep Think completed successfully. Total tokens: {total_tokens}")
+
+ return final_response, total_tokens
def _parse_config(request_config: Dict[str, Any]) -> Dict[str, Any]:
"""Parse and validate configuration parameters."""
diff --git a/optillm/plugins/longcepo/__init__.py b/optillm/plugins/longcepo/__init__.py
index e69de29b..e88bb231 100644
--- a/optillm/plugins/longcepo/__init__.py
+++ b/optillm/plugins/longcepo/__init__.py
@@ -0,0 +1,10 @@
+"""LongCePO Plugin Package
+
+Implementation of Long-Context Cerebras Planning and Optimization method.
+"""
+
+from .main import run_longcepo
+
+__version__ = "1.0.0"
+__author__ = "Cerebras"
+__all__ = ['run_longcepo']
\ No newline at end of file
diff --git a/optillm/plugins/longcepo.py b/optillm/plugins/longcepo_plugin.py
similarity index 100%
rename from optillm/plugins/longcepo.py
rename to optillm/plugins/longcepo_plugin.py
diff --git a/optillm/plugins/spl/__init__.py b/optillm/plugins/spl/__init__.py
index 99df98d3..81c969ec 100644
--- a/optillm/plugins/spl/__init__.py
+++ b/optillm/plugins/spl/__init__.py
@@ -1,3 +1,7 @@
"""
System Prompt Learning (SPL) plugin module initialization.
"""
+
+from .main import run_spl
+
+__all__ = ['run_spl']
diff --git a/optillm/plugins/spl.py b/optillm/plugins/spl_plugin.py
similarity index 100%
rename from optillm/plugins/spl.py
rename to optillm/plugins/spl_plugin.py
diff --git a/optillm/thinkdeeper.py b/optillm/thinkdeeper.py
index e13828c3..321e9520 100644
--- a/optillm/thinkdeeper.py
+++ b/optillm/thinkdeeper.py
@@ -168,8 +168,8 @@ def reasoning_effort(self, messages) -> str:
response = "".join(response_chunks)
full_response = f"{self.config['start_think_token']}\n{self.config['prefill']}{response}"
- logger.debug(f"Final response length: {len(full_response)} chars, Total thoughts: {self.thought_count}")
- return full_response
+ logger.debug(f"Final response length: {len(full_response)} chars, Total thoughts: {self.thought_count}, Thinking tokens: {n_thinking_tokens}")
+ return full_response, n_thinking_tokens
def thinkdeeper_decode(
model: PreTrainedModel,
@@ -192,8 +192,8 @@ def thinkdeeper_decode(
try:
processor = ThinkDeeperProcessor(config, tokenizer, model)
- response = processor.reasoning_effort(messages)
- return response
+ response, reasoning_tokens = processor.reasoning_effort(messages)
+ return response, reasoning_tokens
except Exception as e:
logger.error(f"Error in ThinkDeeper processing: {str(e)}")
diff --git a/optillm/thinkdeeper_mlx.py b/optillm/thinkdeeper_mlx.py
index 043e2876..42c099d6 100644
--- a/optillm/thinkdeeper_mlx.py
+++ b/optillm/thinkdeeper_mlx.py
@@ -243,7 +243,8 @@ def reasoning_effort(self, messages) -> str:
response_content = "".join(response_chunks)
full_response = f"{self.config['start_think_token']}\n{self.config['prefill']}{response_content}"
- return full_response
+ logger.debug(f"MLX Final response length: {len(full_response)} chars, Thinking tokens: {n_thinking_tokens}")
+ return full_response, n_thinking_tokens
def _generate_chunk(self, prompt: str, max_tokens: int, temperature: float) -> str:
"""Generate a small chunk of text using MLX with proper sampler"""
@@ -319,8 +320,8 @@ def thinkdeeper_decode_mlx(
try:
processor = MLXThinkDeeperProcessor(config, tokenizer, model)
- response = processor.reasoning_effort(messages)
- return response
+ response, reasoning_tokens = processor.reasoning_effort(messages)
+ return response, reasoning_tokens
except Exception as e:
logger.error(f"Error in MLX ThinkDeeper processing: {str(e)}")
diff --git a/pyproject.toml b/pyproject.toml
index 6b1d7edc..841cb9d3 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
[project]
name = "optillm"
-version = "0.1.25"
+version = "0.1.26"
description = "An optimizing inference proxy for LLMs."
readme = "README.md"
license = "Apache-2.0"
@@ -48,10 +48,23 @@ dependencies = [
"sentencepiece",
"mcp",
"adaptive-classifier",
+ "datasets",
+ "selenium",
+ "webdriver-manager",
# MLX support for Apple Silicon optimization
'mlx-lm>=0.24.0; platform_machine=="arm64" and sys_platform=="darwin"',
]
+[project.optional-dependencies]
+eval = [
+ "tabulate",
+ "accelerate",
+ "huggingface_hub",
+ "httpx",
+ "tqdm",
+ "pandas",
+]
+
[project.urls]
Homepage = "https://github.com/codelion/optillm"
Repository = "https://github.com/codelion/optillm"
diff --git a/requirements.txt b/requirements.txt
index 9fe9ffe1..7b4be468 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -30,6 +30,7 @@ cerebras_cloud_sdk
outlines[transformers]
sentencepiece
adaptive-classifier
+datasets
mcp
# MLX support for Apple Silicon optimization
mlx-lm>=0.24.0; platform_machine=="arm64" and sys_platform=="darwin"
\ No newline at end of file
diff --git a/scripts/eval_simpleqa_benchmark.py b/scripts/eval_simpleqa_benchmark.py
new file mode 100644
index 00000000..5c653d0e
--- /dev/null
+++ b/scripts/eval_simpleqa_benchmark.py
@@ -0,0 +1,560 @@
+#!/usr/bin/env python3
+"""
+SimpleQA Evaluation Script for OptILLM
+
+This script evaluates OptILLM's effectiveness on the SimpleQA factuality benchmark,
+with support for comparing different models and approaches (baseline, web_search, deep_research).
+
+SimpleQA is OpenAI's factuality benchmark containing 4,326 short, fact-seeking questions
+designed to be challenging for frontier models.
+"""
+
+import argparse
+import json
+import os
+import logging
+import re
+import csv
+import time
+import pandas as pd
+from datetime import datetime
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple, Any
+from tqdm import tqdm
+import requests
+from urllib.parse import urlparse
+import httpx
+from openai import OpenAI
+
+# Setup logging
+logging.basicConfig(
+ level=logging.INFO,
+ format='%(asctime)s - %(levelname)s - %(message)s',
+ handlers=[
+ logging.StreamHandler()
+ ]
+)
+logger = logging.getLogger(__name__)
+
+# Constants
+SIMPLEQA_CSV_URL = "https://openaipublic.blob.core.windows.net/simple-evals/simple_qa_test_set.csv"
+DEFAULT_TIMEOUT = 600 # 10 minutes for potentially long research operations
+DEFAULT_GRADER_MODEL = "gpt-4o"
+DEFAULT_BASE_URL = "http://localhost:8000/v1"
+
+# SimpleQA grading prompt based on OpenAI's implementation
+GRADING_PROMPT = """You are grading a response to a factual question. Your task is to classify the response as:
+
+A) CORRECT - The response contains the correct information and directly answers the question
+B) INCORRECT - The response contains factual errors or contradictory information
+C) NOT_ATTEMPTED - The response doesn't attempt to answer or lacks the required information
+
+Guidelines:
+- Semantic meaning matters more than exact wording
+- Minor variations in names/spelling are acceptable
+- Hedging is allowed if core information is correct
+- Partial answers can be correct if they contain the key fact
+- Only grade as INCORRECT if there are clear factual errors
+
+Question: {question}
+Gold Answer: {gold_answer}
+Response: {response}
+
+Grade (A/B/C):"""
+
+
+def remove_thinking_blocks(text: str) -> str:
+ """Remove ... blocks from the response."""
+ if not text:
+ return text
+
+ if '' in text:
+ # Get everything after the last tag
+ parts = text.split('')
+ return parts[-1].strip()
+ elif '' in text and '' not in text:
+ # Handle truncated responses (no closing tag)
+ parts = text.split('')
+ return parts[0].strip() if len(parts) > 1 and parts[0] else ""
+
+ return text
+
+
+class SimpleQAEvaluator:
+ """Main evaluator class for SimpleQA benchmark"""
+
+ def __init__(self,
+ model: str,
+ approach: str,
+ base_url: str = DEFAULT_BASE_URL,
+ grader_model: str = DEFAULT_GRADER_MODEL,
+ timeout: int = DEFAULT_TIMEOUT,
+ cache_dir: str = "cache",
+ output_dir: str = "results"):
+ self.model = model
+ self.approach = approach
+ self.base_url = base_url
+ self.grader_model = grader_model
+ self.timeout = timeout
+ self.cache_dir = Path(cache_dir)
+ self.output_dir = Path(output_dir)
+
+ # Create directories
+ self.cache_dir.mkdir(exist_ok=True)
+ self.output_dir.mkdir(exist_ok=True)
+
+ # Setup OptILLM client with extended timeout
+ self.optillm_client = OpenAI(
+ api_key="optillm",
+ base_url=base_url,
+ timeout=httpx.Timeout(timeout, connect=5.0),
+ max_retries=0
+ )
+
+ # Setup grader client (use OptILLM for grading)
+ try:
+ self.grader_client = OpenAI(
+ api_key="optillm",
+ base_url=base_url,
+ timeout=httpx.Timeout(timeout, connect=5.0),
+ max_retries=0
+ )
+ logger.info("Using OptILLM for grading responses")
+ except Exception as e:
+ logger.warning(f"Could not initialize grader client: {e}")
+ logger.warning("Grading will be skipped.")
+ self.grader_client = None
+
+ # Results tracking
+ self.results = []
+ self.metrics = {
+ "correct": 0,
+ "incorrect": 0,
+ "not_attempted": 0,
+ "errors": 0,
+ "total_processed": 0
+ }
+
+ def download_dataset(self) -> str:
+ """Download SimpleQA dataset if not cached"""
+ cache_file = self.cache_dir / "simple_qa_test_set.csv"
+
+ if cache_file.exists():
+ logger.info(f"Using cached dataset: {cache_file}")
+ return str(cache_file)
+
+ logger.info(f"Downloading SimpleQA dataset from {SIMPLEQA_CSV_URL}")
+
+ try:
+ response = requests.get(SIMPLEQA_CSV_URL, timeout=30)
+ response.raise_for_status()
+
+ with open(cache_file, 'wb') as f:
+ f.write(response.content)
+
+ logger.info(f"Dataset downloaded to {cache_file}")
+ return str(cache_file)
+
+ except Exception as e:
+ logger.error(f"Failed to download dataset: {e}")
+ raise
+
+ def load_dataset(self, num_samples: Optional[int] = None, start_index: int = 0) -> List[Dict]:
+ """Load and parse SimpleQA dataset"""
+ dataset_file = self.download_dataset()
+
+ questions = []
+
+ try:
+ with open(dataset_file, 'r', encoding='utf-8') as f:
+ reader = csv.DictReader(f)
+
+ for i, row in enumerate(reader):
+ if i < start_index:
+ continue
+
+ if num_samples and len(questions) >= num_samples:
+ break
+
+ # Parse metadata if it's JSON string
+ try:
+ metadata = json.loads(row['metadata']) if row['metadata'] else {}
+ except:
+ metadata = {}
+
+ question_data = {
+ 'id': i,
+ 'metadata': metadata,
+ 'question': row['problem'],
+ 'gold_answer': row['answer']
+ }
+ questions.append(question_data)
+
+ logger.info(f"Loaded {len(questions)} questions from dataset")
+ return questions
+
+ except Exception as e:
+ logger.error(f"Failed to load dataset: {e}")
+ raise
+
+ def get_approach_config(self) -> Dict:
+ """Get configuration for specific approach"""
+ if self.approach == "none":
+ return {}
+ elif self.approach == "web_search":
+ return {
+ "num_results": 10,
+ "headless": True,
+ "timeout": 30
+ }
+ elif self.approach == "deep_research":
+ return {
+ "max_iterations": 1,
+ "max_sources": 10
+ }
+ else:
+ return {}
+
+ def query_optillm(self, question: str) -> Tuple[str, bool]:
+ """Query OptILLM with the specified approach"""
+ try:
+ # Determine model name based on approach
+ if self.approach == "none":
+ model_name = self.model
+ else:
+ model_name = f"{self.approach}-{self.model}"
+
+ # Create messages
+ messages = [
+ {
+ "role": "system",
+ "content": "You are a helpful assistant that provides accurate, factual answers to questions. Be direct and concise."
+ },
+ {
+ "role": "user",
+ "content": question
+ }
+ ]
+
+ # Add approach-specific configuration
+ extra_body = {}
+ approach_config = self.get_approach_config()
+ if approach_config:
+ extra_body.update(approach_config)
+
+ logger.debug(f"Querying model: {model_name}")
+ logger.debug(f"Question: {question}")
+
+ response = self.optillm_client.chat.completions.create(
+ model=model_name,
+ messages=messages,
+ extra_body=extra_body if extra_body else None,
+ max_tokens=4096,
+ temperature=0.6
+ )
+
+ answer = response.choices[0].message.content
+ answer = remove_thinking_blocks(answer)
+ logger.debug(f"Response: {answer}")
+
+ return answer, True
+
+ except Exception as e:
+ logger.error(f"Error querying OptILLM: {e}")
+ return f"Error: {str(e)}", False
+
+ def grade_response(self, question: str, gold_answer: str, response: str) -> str:
+ """Grade response using SimpleQA methodology"""
+ if not self.grader_client:
+ return "NOT_GRADED"
+
+ try:
+ grading_prompt = GRADING_PROMPT.format(
+ question=question,
+ gold_answer=gold_answer,
+ response=response
+ )
+
+ grader_response = self.grader_client.chat.completions.create(
+ model=self.grader_model,
+ messages=[{"role": "user", "content": grading_prompt}],
+ temperature=0.6,
+ max_tokens=4096
+ )
+
+ grade_text = grader_response.choices[0].message.content.strip()
+
+ # Strip tags if present
+ grade_text = re.sub(r'.*?', '', grade_text, flags=re.DOTALL).strip()
+
+ # Extract grade (A/B/C)
+ if grade_text.startswith('A'):
+ return "CORRECT"
+ elif grade_text.startswith('B'):
+ return "INCORRECT"
+ elif grade_text.startswith('C'):
+ return "NOT_ATTEMPTED"
+ else:
+ logger.warning(f"Unexpected grade format: {grade_text}")
+ return "NOT_GRADED"
+
+ except Exception as e:
+ logger.error(f"Error grading response: {e}")
+ return "ERROR_GRADING"
+
+ def evaluate_question(self, question_data: Dict) -> Dict:
+ """Evaluate a single question"""
+ question = question_data['question']
+ gold_answer = question_data['gold_answer']
+
+ # Query OptILLM
+ response, success = self.query_optillm(question)
+
+ result = {
+ 'id': question_data['id'],
+ 'metadata': question_data['metadata'],
+ 'question': question,
+ 'gold_answer': gold_answer,
+ 'response': response,
+ 'success': success,
+ 'timestamp': datetime.now().isoformat()
+ }
+
+ if success:
+ # Grade the response
+ grade = self.grade_response(question, gold_answer, response)
+ result['grade'] = grade
+
+ # Update metrics
+ if grade == "CORRECT":
+ self.metrics["correct"] += 1
+ elif grade == "INCORRECT":
+ self.metrics["incorrect"] += 1
+ elif grade == "NOT_ATTEMPTED":
+ self.metrics["not_attempted"] += 1
+ else:
+ result['grade'] = "ERROR"
+ self.metrics["errors"] += 1
+
+ self.metrics["total_processed"] += 1
+ return result
+
+ def calculate_metrics(self) -> Dict:
+ """Calculate final evaluation metrics"""
+ total = self.metrics["total_processed"]
+ correct = self.metrics["correct"]
+ incorrect = self.metrics["incorrect"]
+ not_attempted = self.metrics["not_attempted"]
+ errors = self.metrics["errors"]
+
+ if total == 0:
+ return {"error": "No questions processed"}
+
+ # Basic percentages
+ accuracy = (correct / total) * 100 if total > 0 else 0
+ attempted = correct + incorrect
+ correct_given_attempted = (correct / attempted) * 100 if attempted > 0 else 0
+
+ # F1 score calculation (treating correct as TP, incorrect as FP, not_attempted as FN)
+ precision = correct / (correct + incorrect) if (correct + incorrect) > 0 else 0
+ recall = correct / (correct + not_attempted) if (correct + not_attempted) > 0 else 0
+ f1_score = 2 * (precision * recall) / (precision + recall) if (precision + recall) > 0 else 0
+
+ return {
+ "total_questions": total,
+ "correct": correct,
+ "incorrect": incorrect,
+ "not_attempted": not_attempted,
+ "errors": errors,
+ "accuracy": accuracy,
+ "correct_given_attempted": correct_given_attempted,
+ "precision": precision,
+ "recall": recall,
+ "f1_score": f1_score,
+ "attempted_rate": (attempted / total) * 100 if total > 0 else 0
+ }
+
+ def save_results(self, timestamp: str) -> Tuple[str, str, str]:
+ """Save evaluation results to files"""
+ # Create output directory for this run
+ run_dir = self.output_dir / f"simpleqa_{self.model}_{self.approach}"
+ run_dir.mkdir(parents=True, exist_ok=True)
+
+ # File paths
+ detailed_file = run_dir / f"{timestamp}_detailed.json"
+ metrics_file = run_dir / f"{timestamp}_metrics.json"
+ summary_file = run_dir / f"{timestamp}_summary.csv"
+
+ # Save detailed results
+ with open(detailed_file, 'w') as f:
+ json.dump(self.results, f, indent=2)
+
+ # Calculate and save metrics
+ final_metrics = self.calculate_metrics()
+ final_metrics.update({
+ "model": self.model,
+ "approach": self.approach,
+ "timestamp": timestamp,
+ "base_url": self.base_url,
+ "grader_model": self.grader_model
+ })
+
+ with open(metrics_file, 'w') as f:
+ json.dump(final_metrics, f, indent=2)
+
+ # Save CSV summary
+ df = pd.DataFrame(self.results)
+ df.to_csv(summary_file, index=False)
+
+ logger.info(f"Results saved to {run_dir}")
+
+ return str(detailed_file), str(metrics_file), str(summary_file)
+
+ def run_evaluation(self,
+ num_samples: Optional[int] = None,
+ start_index: int = 0) -> Dict:
+ """Run the complete evaluation"""
+ timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+
+ logger.info(f"Starting SimpleQA evaluation")
+ logger.info(f"Model: {self.model}")
+ logger.info(f"Approach: {self.approach}")
+ logger.info(f"Base URL: {self.base_url}")
+ logger.info(f"Timeout: {self.timeout}s")
+
+ # Load dataset
+ questions = self.load_dataset(num_samples, start_index)
+
+ # Run evaluation with progress bar
+ for question_data in tqdm(questions, desc="Evaluating questions"):
+ try:
+ result = self.evaluate_question(question_data)
+ self.results.append(result)
+
+ # Log progress periodically
+ if len(self.results) % 10 == 0:
+ metrics = self.calculate_metrics()
+ logger.info(f"Progress: {len(self.results)}/{len(questions)} - "
+ f"Accuracy: {metrics['accuracy']:.1f}%")
+
+ except KeyboardInterrupt:
+ logger.info("Evaluation interrupted by user")
+ break
+ except Exception as e:
+ logger.error(f"Error evaluating question {question_data['id']}: {e}")
+ continue
+
+ # Save results
+ detailed_file, metrics_file, summary_file = self.save_results(timestamp)
+
+ # Calculate final metrics
+ final_metrics = self.calculate_metrics()
+
+ logger.info("Evaluation completed!")
+ logger.info(f"Total questions: {final_metrics['total_questions']}")
+ logger.info(f"Accuracy: {final_metrics['accuracy']:.1f}%")
+ logger.info(f"F1 Score: {final_metrics['f1_score']:.3f}")
+ logger.info(f"Correct: {final_metrics['correct']}")
+ logger.info(f"Incorrect: {final_metrics['incorrect']}")
+ logger.info(f"Not Attempted: {final_metrics['not_attempted']}")
+
+ return final_metrics
+
+
+def parse_args():
+ """Parse command line arguments"""
+ parser = argparse.ArgumentParser(
+ description="Evaluate OptILLM on SimpleQA factuality benchmark"
+ )
+
+ # Model and approach
+ parser.add_argument("--model", type=str, default="gpt-4o-mini",
+ help="Model to evaluate (default: gpt-4o-mini)")
+ parser.add_argument("--approach", type=str, default="none",
+ choices=["none", "web_search", "deep_research"],
+ help="Approach to use (default: none)")
+
+ # Server configuration
+ parser.add_argument("--base-url", type=str, default=DEFAULT_BASE_URL,
+ help=f"OptILLM base URL (default: {DEFAULT_BASE_URL})")
+ parser.add_argument("--timeout", type=int, default=DEFAULT_TIMEOUT,
+ help=f"Request timeout in seconds (default: {DEFAULT_TIMEOUT})")
+
+ # Grading configuration
+ parser.add_argument("--grader-model", type=str, default=DEFAULT_GRADER_MODEL,
+ help=f"Model for grading responses (default: {DEFAULT_GRADER_MODEL})")
+
+ # Evaluation parameters
+ parser.add_argument("--num-samples", type=int, default=None,
+ help="Number of questions to evaluate (default: all)")
+ parser.add_argument("--start-index", type=int, default=0,
+ help="Start from specific question index (default: 0)")
+
+ # Search-specific parameters
+ parser.add_argument("--num-search-results", type=int, default=10,
+ help="Number of search results per query (default: 10)")
+ parser.add_argument("--headless", action="store_true",
+ help="Run browser in headless mode for web search")
+
+ # Output configuration
+ parser.add_argument("--cache-dir", type=str, default="cache",
+ help="Directory for caching dataset (default: cache)")
+ parser.add_argument("--output-dir", type=str, default="results",
+ help="Directory for saving results (default: results)")
+
+ # Debugging
+ parser.add_argument("--verbose", action="store_true",
+ help="Enable verbose logging")
+
+ return parser.parse_args()
+
+
+def main():
+ """Main entry point"""
+ args = parse_args()
+
+ if args.verbose:
+ logging.getLogger().setLevel(logging.DEBUG)
+
+ # Create evaluator
+ evaluator = SimpleQAEvaluator(
+ model=args.model,
+ approach=args.approach,
+ base_url=args.base_url,
+ grader_model=args.grader_model,
+ timeout=args.timeout,
+ cache_dir=args.cache_dir,
+ output_dir=args.output_dir
+ )
+
+ try:
+ # Run evaluation
+ metrics = evaluator.run_evaluation(
+ num_samples=args.num_samples,
+ start_index=args.start_index
+ )
+
+ print("\n" + "="*50)
+ print("EVALUATION SUMMARY")
+ print("="*50)
+ print(f"Model: {args.model}")
+ print(f"Approach: {args.approach}")
+ print(f"Questions: {metrics['total_questions']}")
+ print(f"Accuracy: {metrics['accuracy']:.1f}%")
+ print(f"F1 Score: {metrics['f1_score']:.3f}")
+ print(f"Correct: {metrics['correct']}")
+ print(f"Incorrect: {metrics['incorrect']}")
+ print(f"Not Attempted: {metrics['not_attempted']}")
+
+ if metrics['errors'] > 0:
+ print(f"Errors: {metrics['errors']}")
+
+ except KeyboardInterrupt:
+ print("\nEvaluation interrupted by user")
+ except Exception as e:
+ logger.error(f"Evaluation failed: {e}")
+ raise
+
+
+if __name__ == "__main__":
+ main()
\ No newline at end of file
diff --git a/scripts/requirements.txt b/scripts/requirements.txt
index bcc1e07f..d850c8cc 100644
--- a/scripts/requirements.txt
+++ b/scripts/requirements.txt
@@ -2,3 +2,8 @@ tabulate
datasets
accelerate
huggingface_hub
+openai
+httpx
+tqdm
+requests
+pandas
diff --git a/tests/test.py b/tests/test.py
index 62989d41..5269d695 100644
--- a/tests/test.py
+++ b/tests/test.py
@@ -30,8 +30,8 @@
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)
-# OpenAI API configuration
-API_KEY = os.environ.get("OPENAI_API_KEY")
+# API configuration - default to local inference for testing
+API_KEY = os.environ.get("OPENAI_API_KEY", "optillm")
# Mock OpenAI client for testing purposes
class MockOpenAIClient:
@@ -150,14 +150,23 @@ def main():
args.test_cases = os.path.join(script_dir, "test_cases.json")
# If using local inference mode, override model to a local model
- if os.environ.get("OPTILLM_API_KEY") == "optillm" and args.model == "gpt-4o-mini":
+ if API_KEY == "optillm" and args.model == "gpt-4o-mini":
args.model = "Qwen/Qwen2.5-0.5B-Instruct"
logger.info(f"Using local model: {args.model}")
+
+ # Set environment variable for local inference
+ if API_KEY == "optillm":
+ os.environ["OPTILLM_API_KEY"] = "optillm"
test_cases = load_test_cases(args.test_cases)
+ # Use local inference by default for testing
if args.base_url:
client = OpenAI(api_key=API_KEY, base_url=args.base_url)
+ elif API_KEY == "optillm":
+ # Use local inference endpoint
+ client = OpenAI(api_key=API_KEY, base_url="http://localhost:8000/v1")
+ logger.info("Using local inference endpoint: http://localhost:8000/v1")
else:
client = OpenAI(api_key=API_KEY)
# client = LiteLLMWrapper()
diff --git a/tests/test_api_compatibility.py b/tests/test_api_compatibility.py
index e33d6e92..7a11a8ae 100644
--- a/tests/test_api_compatibility.py
+++ b/tests/test_api_compatibility.py
@@ -103,6 +103,98 @@ def test_streaming(client):
assert len(content_chunks) > 0
+def test_reasoning_tokens_in_response(client):
+ """Test that reasoning tokens are included in API responses"""
+ response = client.chat.completions.create(
+ model="gpt-4o-mini",
+ messages=[
+ {"role": "system", "content": "Think step by step and show your reasoning."},
+ {"role": "user", "content": "What is 15 × 23? Please think through this step by step."}
+ ],
+ max_tokens=100
+ )
+
+ # Check basic response structure
+ assert hasattr(response, 'choices')
+ assert len(response.choices) > 0
+ assert hasattr(response, 'usage')
+
+ # Check that completion_tokens_details exists and has reasoning_tokens
+ assert hasattr(response.usage, 'completion_tokens_details')
+ assert hasattr(response.usage.completion_tokens_details, 'reasoning_tokens')
+
+ # reasoning_tokens should be an integer >= 0
+ reasoning_tokens = response.usage.completion_tokens_details.reasoning_tokens
+ assert isinstance(reasoning_tokens, int)
+ assert reasoning_tokens >= 0
+
+
+def test_reasoning_tokens_with_thinking_prompt(client):
+ """Test reasoning tokens with a prompt designed to trigger thinking"""
+ response = client.chat.completions.create(
+ model="gpt-4o-mini",
+ messages=[
+ {"role": "system", "content": "You are a helpful assistant. Use tags to show your reasoning process."},
+ {"role": "user", "content": "I have 12 apples. I eat 3, give away 4, and buy 7 more. How many apples do I have now?"}
+ ],
+ max_tokens=150
+ )
+
+ # Basic checks
+ assert hasattr(response, 'usage')
+ assert hasattr(response.usage, 'completion_tokens_details')
+ assert hasattr(response.usage.completion_tokens_details, 'reasoning_tokens')
+
+ reasoning_tokens = response.usage.completion_tokens_details.reasoning_tokens
+ assert isinstance(reasoning_tokens, int)
+ assert reasoning_tokens >= 0
+
+ # If the model used thinking tags, reasoning_tokens should be > 0
+ # (This depends on the model's response, so we just check the structure)
+
+
+def test_reasoning_tokens_with_multiple_responses(client):
+ """Test reasoning tokens with n > 1"""
+ response = client.chat.completions.create(
+ model="gpt-4o-mini",
+ messages=[
+ {"role": "user", "content": "Think about this: What's 2+2?"}
+ ],
+ n=2,
+ max_tokens=50
+ )
+
+ # Should have 2 choices
+ assert len(response.choices) == 2
+
+ # Should have reasoning token information
+ assert hasattr(response.usage, 'completion_tokens_details')
+ assert hasattr(response.usage.completion_tokens_details, 'reasoning_tokens')
+
+ reasoning_tokens = response.usage.completion_tokens_details.reasoning_tokens
+ assert isinstance(reasoning_tokens, int)
+ assert reasoning_tokens >= 0
+
+
+def test_reasoning_tokens_backward_compatibility(client):
+ """Test that responses without thinking still work normally"""
+ response = client.chat.completions.create(
+ model="gpt-4o-mini",
+ messages=[
+ {"role": "user", "content": "Say hello"}
+ ],
+ max_tokens=10
+ )
+
+ # Should still have reasoning token structure, but with 0 tokens
+ assert hasattr(response.usage, 'completion_tokens_details')
+ assert hasattr(response.usage.completion_tokens_details, 'reasoning_tokens')
+
+ reasoning_tokens = response.usage.completion_tokens_details.reasoning_tokens
+ assert isinstance(reasoning_tokens, int)
+ assert reasoning_tokens >= 0 # Usually 0 for simple responses
+
+
if __name__ == "__main__":
# Run basic tests if pytest not available
client = OpenAI(
@@ -110,24 +202,39 @@ def test_streaming(client):
base_url="http://localhost:8000/v1"
)
- print("Running basic API compatibility tests...")
+ print("Running API compatibility tests...")
+
+ tests = [
+ ("Basic completion", test_basic_completion),
+ ("N parameter", test_n_parameter),
+ ("Approach prefix", test_approach_prefix),
+ ("Extra body approach", test_extra_body_approach),
+ ("Streaming", test_streaming),
+ ("Reasoning tokens in response", test_reasoning_tokens_in_response),
+ ("Reasoning tokens with thinking prompt", test_reasoning_tokens_with_thinking_prompt),
+ ("Reasoning tokens with multiple responses", test_reasoning_tokens_with_multiple_responses),
+ ("Reasoning tokens backward compatibility", test_reasoning_tokens_backward_compatibility),
+ ]
- try:
- test_basic_completion(client)
- print("✅ Basic completion test passed")
- except Exception as e:
- print(f"❌ Basic completion test failed: {e}")
+ passed = 0
+ failed = 0
- try:
- test_n_parameter(client)
- print("✅ N parameter test passed")
- except Exception as e:
- print(f"❌ N parameter test failed: {e}")
+ for test_name, test_func in tests:
+ try:
+ print(f"Running {test_name}...", end=' ')
+ test_func(client)
+ print("✅ PASSED")
+ passed += 1
+ except Exception as e:
+ print(f"❌ FAILED: {e}")
+ failed += 1
- try:
- test_approach_prefix(client)
- print("✅ Approach prefix test passed")
- except Exception as e:
- print(f"❌ Approach prefix test failed: {e}")
+ print(f"\n=== Test Summary ===")
+ print(f"Passed: {passed}")
+ print(f"Failed: {failed}")
+ print(f"Total: {passed + failed}")
- print("\nDone!")
\ No newline at end of file
+ if failed == 0:
+ print("🎉 All tests passed!")
+ else:
+ print(f"⚠️ {failed} test(s) failed.")
\ No newline at end of file
diff --git a/tests/test_approaches.py b/tests/test_approaches.py
index 10ea67f9..1749a301 100644
--- a/tests/test_approaches.py
+++ b/tests/test_approaches.py
@@ -4,7 +4,6 @@
Tests the basic structure of approaches without requiring actual model inference
"""
-import pytest
import sys
import os
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
diff --git a/tests/test_cases.json b/tests/test_cases.json
index 7b10ae43..d136ca3b 100644
--- a/tests/test_cases.json
+++ b/tests/test_cases.json
@@ -43,5 +43,30 @@
"name": "Simple Math Problem",
"system_prompt": "You are a helpful assistant.",
"query": "What is 2 + 2?"
+ },
+ {
+ "name": "Reasoning Token Test - Complex Logic",
+ "system_prompt": "You are an AI assistant that thinks step by step. Use tags to show your reasoning process.",
+ "query": "Three friends Alice, Bob, and Charlie each have a different number of marbles. Alice has twice as many as Bob. Charlie has 3 more than Alice. Together they have 23 marbles. How many marbles does each person have?"
+ },
+ {
+ "name": "Reasoning Token Test - Strategic Thinking",
+ "system_prompt": "Think carefully before responding. Show your work using thinking tags.",
+ "query": "You're playing a game where you can choose door A or door B. Behind one door is a prize worth $1000, behind the other is nothing. You know that if the prize is behind door A, there's a 70% chance a light above door A will flash. If the prize is behind door B, there's a 30% chance the light above door A will flash. The light above door A is flashing. Which door should you choose?"
+ },
+ {
+ "name": "Reasoning Token Test - Multi-Step Problem",
+ "system_prompt": "Please think through this problem step by step, showing your reasoning.",
+ "query": "A bakery sells cupcakes in boxes of 6 and cookies in boxes of 8. If someone buys the same number of cupcakes and cookies, what is the smallest number of each type of baked good they could buy? Show all your work."
+ },
+ {
+ "name": "Reasoning Token Test - Counter-intuitive",
+ "system_prompt": "This problem might seem simple but requires careful analysis. Think it through.",
+ "query": "In a family with two children, you know that at least one of them is a boy. What is the probability that both children are boys? Explain your reasoning carefully."
+ },
+ {
+ "name": "Reasoning Token Test - Algorithm Design",
+ "system_prompt": "Think through the algorithm design process step by step.",
+ "query": "Design an efficient algorithm to find the second largest element in an unsorted array. Explain your approach, analyze the time complexity, and provide pseudocode."
}
]
diff --git a/tests/test_ci_quick.py b/tests/test_ci_quick.py
index 332ae409..7b90e8d8 100644
--- a/tests/test_ci_quick.py
+++ b/tests/test_ci_quick.py
@@ -34,9 +34,25 @@
import optillm.plugins.privacy_plugin
import optillm.plugins.genselect_plugin
import optillm.plugins.majority_voting_plugin
- print("✅ Plugin modules exist and can be imported")
+ print("✅ Basic plugin modules exist and can be imported")
except Exception as e:
- print(f"❌ Plugin import test failed: {e}")
+ print(f"❌ Basic plugin import test failed: {e}")
+
+# Test plugin subdirectory imports (critical for issue #220)
+try:
+ from optillm.plugins.deepthink import SelfDiscover, UncertaintyRoutedCoT
+ from optillm.plugins.deep_research import DeepResearcher
+ from optillm.plugins.longcepo import run_longcepo
+ from optillm.plugins.spl import run_spl
+ print("✅ Plugin submodule imports working - no relative import errors")
+except ImportError as e:
+ if "attempted relative import" in str(e):
+ print(f"❌ Critical: Relative import error detected: {e}")
+ sys.exit(1)
+ else:
+ print(f"❌ Plugin submodule import error: {e}")
+except Exception as e:
+ print(f"❌ Plugin submodule import error: {e}")
# Test approach parsing
try:
diff --git a/tests/test_plugins.py b/tests/test_plugins.py
index 41fbfd2b..63f0fc12 100644
--- a/tests/test_plugins.py
+++ b/tests/test_plugins.py
@@ -27,7 +27,10 @@ def test_plugin_module_imports():
'optillm.plugins.genselect_plugin',
'optillm.plugins.majority_voting_plugin',
'optillm.plugins.web_search_plugin',
- 'optillm.plugins.deep_research_plugin'
+ 'optillm.plugins.deep_research_plugin',
+ 'optillm.plugins.deepthink_plugin',
+ 'optillm.plugins.longcepo_plugin',
+ 'optillm.plugins.spl_plugin'
]
for module_name in plugin_modules:
@@ -48,7 +51,7 @@ def test_plugin_approach_detection():
load_plugins()
# Check if known plugins are loaded
- expected_plugins = ["memory", "readurls", "privacy", "web_search", "deep_research"]
+ expected_plugins = ["memory", "readurls", "privacy", "web_search", "deep_research", "deepthink", "longcepo", "spl"]
for plugin_name in expected_plugins:
assert plugin_name in plugin_approaches, f"Plugin {plugin_name} not loaded"
@@ -76,8 +79,8 @@ def test_majority_voting_plugin():
import optillm.plugins.majority_voting_plugin as plugin
assert hasattr(plugin, 'run')
assert hasattr(plugin, 'SLUG')
- assert hasattr(plugin, 'extract_answer')
- assert hasattr(plugin, 'normalize_answer')
+ assert hasattr(plugin, 'extract_final_answer')
+ assert hasattr(plugin, 'normalize_response')
assert plugin.SLUG == "majority_voting"
@@ -100,6 +103,96 @@ def test_deep_research_plugin():
assert plugin.SLUG == "deep_research"
+def test_deepthink_plugin_imports():
+ """Test deepthink plugin and its submodules can be imported"""
+ # Test main plugin
+ import optillm.plugins.deepthink_plugin as plugin
+ assert hasattr(plugin, 'run')
+ assert hasattr(plugin, 'SLUG')
+ assert plugin.SLUG == "deepthink"
+
+ # Test submodules can be imported
+ from optillm.plugins.deepthink import SelfDiscover, UncertaintyRoutedCoT
+ assert SelfDiscover is not None
+ assert UncertaintyRoutedCoT is not None
+
+
+def test_longcepo_plugin():
+ """Test longcepo plugin module"""
+ import optillm.plugins.longcepo_plugin as plugin
+ assert hasattr(plugin, 'run')
+ assert hasattr(plugin, 'SLUG')
+ assert plugin.SLUG == "longcepo"
+
+ # Test submodule can be imported
+ from optillm.plugins.longcepo import run_longcepo
+ assert run_longcepo is not None
+
+
+def test_spl_plugin():
+ """Test spl plugin module"""
+ import optillm.plugins.spl_plugin as plugin
+ assert hasattr(plugin, 'run')
+ assert hasattr(plugin, 'SLUG')
+ assert plugin.SLUG == "spl"
+
+ # Test submodule can be imported
+ from optillm.plugins.spl import run_spl
+ assert run_spl is not None
+
+
+def test_plugin_subdirectory_imports():
+ """Test all plugins with subdirectories can import their submodules"""
+ # Test deep_research
+ from optillm.plugins.deep_research import DeepResearcher
+ assert DeepResearcher is not None
+
+ # Test deepthink
+ from optillm.plugins.deepthink import SelfDiscover, UncertaintyRoutedCoT
+ assert SelfDiscover is not None
+ assert UncertaintyRoutedCoT is not None
+
+ # Test longcepo
+ from optillm.plugins.longcepo import run_longcepo
+ assert run_longcepo is not None
+
+ # Test spl
+ from optillm.plugins.spl import run_spl
+ assert run_spl is not None
+
+
+def test_no_relative_import_errors():
+ """Test that plugins load without relative import errors"""
+ import importlib
+ import sys
+
+ plugins_with_subdirs = [
+ 'optillm.plugins.deepthink_plugin',
+ 'optillm.plugins.deep_research_plugin',
+ 'optillm.plugins.longcepo_plugin',
+ 'optillm.plugins.spl_plugin'
+ ]
+
+ for plugin_name in plugins_with_subdirs:
+ # Clear any previously loaded modules to test fresh import
+ modules_to_clear = [k for k in sys.modules.keys() if k.startswith(plugin_name)]
+ for mod in modules_to_clear:
+ del sys.modules[mod]
+
+ try:
+ module = importlib.import_module(plugin_name)
+ # Try to access the run function to ensure full initialization works
+ assert hasattr(module, 'run'), f"{plugin_name} missing run function"
+ except ImportError as e:
+ if "attempted relative import" in str(e):
+ if pytest:
+ pytest.fail(f"Relative import error in {plugin_name}: {e}")
+ else:
+ raise AssertionError(f"Relative import error in {plugin_name}: {e}")
+ else:
+ raise
+
+
if __name__ == "__main__":
print("Running plugin tests...")
@@ -145,4 +238,34 @@ def test_deep_research_plugin():
except Exception as e:
print(f"❌ Deep research plugin test failed: {e}")
+ try:
+ test_deepthink_plugin_imports()
+ print("✅ Deepthink plugin imports test passed")
+ except Exception as e:
+ print(f"❌ Deepthink plugin imports test failed: {e}")
+
+ try:
+ test_longcepo_plugin()
+ print("✅ LongCePO plugin test passed")
+ except Exception as e:
+ print(f"❌ LongCePO plugin test failed: {e}")
+
+ try:
+ test_spl_plugin()
+ print("✅ SPL plugin test passed")
+ except Exception as e:
+ print(f"❌ SPL plugin test failed: {e}")
+
+ try:
+ test_plugin_subdirectory_imports()
+ print("✅ Plugin subdirectory imports test passed")
+ except Exception as e:
+ print(f"❌ Plugin subdirectory imports test failed: {e}")
+
+ try:
+ test_no_relative_import_errors()
+ print("✅ No relative import errors test passed")
+ except Exception as e:
+ print(f"❌ Relative import errors test failed: {e}")
+
print("\nDone!")
\ No newline at end of file
diff --git a/tests/test_reasoning_integration.py b/tests/test_reasoning_integration.py
new file mode 100644
index 00000000..b01f871f
--- /dev/null
+++ b/tests/test_reasoning_integration.py
@@ -0,0 +1,318 @@
+#!/usr/bin/env python3
+"""
+Integration tests for reasoning token functionality
+Tests end-to-end integration with approaches that generate thinking
+"""
+
+import sys
+import os
+import unittest
+from unittest.mock import Mock, patch, MagicMock
+import re
+
+# Add parent directory to path for imports
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+# Import the thinkdeeper functions for testing
+from optillm.thinkdeeper import thinkdeeper_decode
+from optillm.thinkdeeper_mlx import thinkdeeper_decode_mlx
+
+
+class MockTokenizer:
+ """Mock tokenizer for testing"""
+ def encode(self, text):
+ # Simple word-based tokenization for testing
+ return text.split()
+
+ def decode(self, tokens):
+ return " ".join(str(t) for t in tokens)
+
+ def apply_chat_template(self, messages, **kwargs):
+ # Simple template that just concatenates messages
+ text = " ".join(msg["content"] for msg in messages)
+ return [[1, 2, 3] + self.encode(text)] # Mock token tensor format
+
+
+class MockModel:
+ """Mock model for testing"""
+ def __init__(self):
+ self.device = "cpu"
+ self.config = Mock()
+ self.generation_config = Mock()
+
+ def __call__(self, **kwargs):
+ # Mock model output with logits
+ class MockOutput:
+ def __init__(self):
+ # Create mock logits tensor
+ import torch
+ self.logits = torch.randn(1, 1, 1000) # batch_size=1, seq_len=1, vocab_size=1000
+
+ return MockOutput()
+
+
+class TestThinkDeeperReasoningTokens(unittest.TestCase):
+ """Test ThinkDeeper approaches return reasoning tokens"""
+
+ def setUp(self):
+ """Set up test fixtures"""
+ self.mock_tokenizer = MockTokenizer()
+ self.mock_model = MockModel()
+ self.test_messages = [
+ {"role": "user", "content": "What is 2 + 2?"}
+ ]
+
+ def test_thinkdeeper_returns_reasoning_tokens(self):
+ """Test that thinkdeeper_decode returns reasoning tokens"""
+ try:
+ # Mock torch operations to avoid actual model inference
+ with patch('torch.tensor') as mock_tensor, \
+ patch('torch.randn') as mock_randn, \
+ patch('torch.multinomial') as mock_multinomial:
+
+ # Set up mocks
+ mock_tensor.return_value = Mock()
+ mock_tensor.return_value.to.return_value = Mock()
+ mock_randn.return_value = Mock()
+ mock_multinomial.return_value = Mock()
+ mock_multinomial.return_value.item.return_value = 50 # Mock token ID for
+
+ # Mock the tokenizer's encode method to return specific tokens
+ def mock_encode(text):
+ if "" in text:
+ return [50] # Token ID for
+ return [1, 2, 3, 4, 5] # Other tokens
+
+ self.mock_tokenizer.encode = mock_encode
+
+ # Mock the model to stop generation quickly
+ generation_count = 0
+ def mock_model_call(**kwargs):
+ nonlocal generation_count
+ generation_count += 1
+
+ class MockOutput:
+ def __init__(self):
+ import torch
+ # After a few calls, return the end think token
+ if generation_count > 3:
+ self.logits = torch.zeros(1, 1, 1000)
+ self.logits[0, 0, 50] = 100 # High logit for end think token
+ else:
+ self.logits = torch.randn(1, 1, 1000)
+
+ return MockOutput()
+
+ self.mock_model.__call__ = mock_model_call
+
+ # Test thinkdeeper_decode
+ result = thinkdeeper_decode(
+ self.mock_model,
+ self.mock_tokenizer,
+ self.test_messages
+ )
+
+ # Should return tuple with (response, reasoning_tokens)
+ self.assertIsInstance(result, tuple)
+ self.assertEqual(len(result), 2)
+
+ response, reasoning_tokens = result
+ self.assertIsInstance(response, str)
+ self.assertIsInstance(reasoning_tokens, int)
+ self.assertGreaterEqual(reasoning_tokens, 0)
+
+ except Exception as e:
+ # If actual thinkdeeper fails due to mocking complexity,
+ # at least verify the function signature changed
+ self.assertIn("too many values to unpack", str(e))
+
+ def test_thinkdeeper_mlx_returns_reasoning_tokens(self):
+ """Test that thinkdeeper_decode_mlx returns reasoning tokens"""
+ try:
+ # Mock MLX operations
+ with patch('mlx.core.array') as mock_array, \
+ patch('mlx.nn.sample') as mock_sample:
+
+ # Set up MLX mocks
+ mock_array.return_value = Mock()
+ mock_sample.return_value = Mock()
+ mock_sample.return_value.item.return_value = 50 # Mock token
+
+ # Mock the model to have MLX-like interface
+ class MockMLXModel:
+ def __call__(self, inputs):
+ # Return mock logits
+ return Mock()
+
+ mlx_model = MockMLXModel()
+
+ # Test thinkdeeper_decode_mlx
+ result = thinkdeeper_decode_mlx(
+ mlx_model,
+ self.mock_tokenizer,
+ self.test_messages
+ )
+
+ # Should return tuple with (response, reasoning_tokens)
+ self.assertIsInstance(result, tuple)
+ self.assertEqual(len(result), 2)
+
+ response, reasoning_tokens = result
+ self.assertIsInstance(response, str)
+ self.assertIsInstance(reasoning_tokens, int)
+ self.assertGreaterEqual(reasoning_tokens, 0)
+
+ except Exception as e:
+ # If actual MLX thinkdeeper fails due to import or mocking,
+ # at least verify the function signature changed
+ if "mlx" not in str(e).lower():
+ self.assertIn("too many values to unpack", str(e))
+
+
+class TestInferenceIntegration(unittest.TestCase):
+ """Test integration with inference.py module"""
+
+ def test_inference_usage_includes_reasoning_tokens(self):
+ """Test that ChatCompletionUsage includes reasoning_tokens"""
+ from optillm.inference import ChatCompletionUsage
+
+ # Test creating usage with reasoning tokens
+ usage = ChatCompletionUsage(
+ prompt_tokens=10,
+ completion_tokens=20,
+ total_tokens=30,
+ reasoning_tokens=5
+ )
+
+ self.assertEqual(usage.prompt_tokens, 10)
+ self.assertEqual(usage.completion_tokens, 20)
+ self.assertEqual(usage.total_tokens, 30)
+ self.assertEqual(usage.reasoning_tokens, 5)
+
+ def test_inference_usage_default_reasoning_tokens(self):
+ """Test that ChatCompletionUsage defaults reasoning_tokens to 0"""
+ from optillm.inference import ChatCompletionUsage
+
+ # Test creating usage without reasoning tokens
+ usage = ChatCompletionUsage(
+ prompt_tokens=10,
+ completion_tokens=20,
+ total_tokens=30
+ )
+
+ self.assertEqual(usage.reasoning_tokens, 0)
+
+ def test_chat_completion_model_dump_includes_reasoning_tokens(self):
+ """Test that ChatCompletion.model_dump includes reasoning_tokens in usage"""
+ from optillm.inference import ChatCompletion
+
+ # Create mock response with reasoning tokens
+ response_dict = {
+ "id": "test-id",
+ "object": "chat.completion",
+ "created": 1234567890,
+ "model": "test-model",
+ "choices": [
+ {
+ "index": 0,
+ "message": {
+ "role": "assistant",
+ "content": "reasoninganswer"
+ },
+ "finish_reason": "stop"
+ }
+ ],
+ "usage": {
+ "prompt_tokens": 10,
+ "completion_tokens": 20,
+ "total_tokens": 30,
+ "reasoning_tokens": 5
+ }
+ }
+
+ completion = ChatCompletion(response_dict)
+ result = completion.model_dump()
+
+ # Check that model_dump includes reasoning_tokens
+ self.assertIn("usage", result)
+ self.assertIn("completion_tokens_details", result["usage"])
+ self.assertIn("reasoning_tokens", result["usage"]["completion_tokens_details"])
+ self.assertEqual(result["usage"]["completion_tokens_details"]["reasoning_tokens"], 5)
+
+
+class TestEndToEndIntegration(unittest.TestCase):
+ """Test end-to-end integration with mocked dependencies"""
+
+ @patch('optillm.get_config')
+ def test_thinkdeeper_approach_with_reasoning_tokens(self, mock_get_config):
+ """Test end-to-end with thinkdeeper approach"""
+ import optillm
+
+ # Set up server config for thinkdeeper
+ optillm.server_config['approach'] = 'none' # Use none to avoid plugin loading issues
+
+ # Mock the OpenAI client to return think tags
+ mock_client = Mock()
+ mock_response = Mock()
+ mock_response.choices = [Mock()]
+ mock_response.choices[0].message.content = "I need to calculate 2+2. Let me think step by step.The answer is 4."
+ mock_response.usage.completion_tokens = 25
+ mock_response.usage.prompt_tokens = 8
+ mock_response.usage.total_tokens = 33
+
+ mock_client.chat.completions.create.return_value = mock_response
+ mock_get_config.return_value = (mock_client, "test-key")
+
+ # Create test client
+ app = optillm.app
+ app.config['TESTING'] = True
+ client = app.test_client()
+
+ # Make request
+ response = client.post('/v1/chat/completions',
+ json={
+ "model": "gpt-4o-mini",
+ "messages": [{"role": "user", "content": "What is 2+2?"}]
+ },
+ headers={"Authorization": "Bearer test-key"})
+
+ self.assertEqual(response.status_code, 200)
+
+ # Check that response includes reasoning tokens
+ data = response.get_json()
+ self.assertIn('usage', data)
+ self.assertIn('completion_tokens_details', data['usage'])
+ self.assertIn('reasoning_tokens', data['usage']['completion_tokens_details'])
+
+ # Should have detected reasoning tokens from the think tags
+ reasoning_tokens = data['usage']['completion_tokens_details']['reasoning_tokens']
+ self.assertGreater(reasoning_tokens, 0)
+ self.assertLess(reasoning_tokens, data['usage']['completion_tokens'])
+
+
+class TestLocalInferenceReasoningTokens(unittest.TestCase):
+ """Test reasoning tokens with local inference if available"""
+
+ def test_local_inference_reasoning_calculation(self):
+ """Test that local inference calculates reasoning tokens correctly"""
+ try:
+ from optillm.inference import InferenceClient
+
+ # Create mock inference client
+ client = InferenceClient()
+
+ # This test mainly verifies the structure exists
+ # Actual inference testing would require models to be available
+ self.assertTrue(hasattr(client, 'chat'))
+
+ except ImportError:
+ # If inference dependencies aren't available, skip
+ self.skipTest("Local inference dependencies not available")
+ except Exception as e:
+ # If other errors occur during initialization, that's still informative
+ self.assertTrue(True, f"InferenceClient initialization: {e}")
+
+
+if __name__ == '__main__':
+ # Run the tests
+ unittest.main(verbosity=2)
\ No newline at end of file
diff --git a/tests/test_reasoning_simple.py b/tests/test_reasoning_simple.py
new file mode 100644
index 00000000..f35e39cf
--- /dev/null
+++ b/tests/test_reasoning_simple.py
@@ -0,0 +1,210 @@
+#!/usr/bin/env python3
+"""
+Simple tests for reasoning token functionality
+Focuses on unit tests that don't require complex mocking
+"""
+
+import sys
+import os
+import unittest
+
+# Add parent directory to path for imports
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+# Import the functions we want to test
+from optillm import count_reasoning_tokens as optillm_count
+from optillm.inference import count_reasoning_tokens as inference_count
+
+
+class TestReasoningTokensCore(unittest.TestCase):
+ """Test core reasoning token functionality"""
+
+ def test_count_reasoning_tokens_with_think_tags(self):
+ """Test counting tokens in think tags"""
+ text = "Let me think about this problem step by stepThe answer is 42"
+
+ result1 = optillm_count(text)
+ result2 = inference_count(text)
+
+ self.assertGreater(result1, 0, "Should count tokens in think tags")
+ self.assertEqual(result1, result2, "Both functions should return same result")
+
+ def test_count_reasoning_tokens_without_think_tags(self):
+ """Test with text that has no think tags"""
+ text = "This is just a regular response without any thinking"
+
+ result1 = optillm_count(text)
+ result2 = inference_count(text)
+
+ self.assertEqual(result1, 0, "Should return 0 for text without think tags")
+ self.assertEqual(result2, 0, "Should return 0 for text without think tags")
+
+ def test_count_reasoning_tokens_multiple_blocks(self):
+ """Test with multiple think tag blocks"""
+ text = """
+ First block of reasoning
+ Some output here
+ Second block with more reasoning
+ Final answer
+ """
+
+ result = optillm_count(text)
+ self.assertGreater(result, 0, "Should count tokens from multiple blocks")
+
+ def test_count_reasoning_tokens_empty_cases(self):
+ """Test edge cases with empty or invalid input"""
+ test_cases = ["", None, 123, ""]
+
+ for case in test_cases:
+ result1 = optillm_count(case)
+ result2 = inference_count(case)
+
+ self.assertGreaterEqual(result1, 0, f"Should handle {case} gracefully")
+ self.assertGreaterEqual(result2, 0, f"Should handle {case} gracefully")
+
+ def test_count_reasoning_tokens_with_mock_tokenizer(self):
+ """Test with a simple mock tokenizer"""
+ class MockTokenizer:
+ def encode(self, text):
+ return text.split() # Simple word-based tokenization
+
+ tokenizer = MockTokenizer()
+ text = "hello world testanswer"
+
+ result = optillm_count(text, tokenizer)
+ self.assertEqual(result, 3, "Should use tokenizer when provided")
+
+ def test_reasoning_tokens_fallback_estimation(self):
+ """Test fallback estimation when tokenizer fails"""
+ class FailingTokenizer:
+ def encode(self, text):
+ raise Exception("Tokenizer failed")
+
+ tokenizer = FailingTokenizer()
+ text = "some reasoning content hereanswer"
+
+ result = optillm_count(text, tokenizer)
+ self.assertGreater(result, 0, "Should fallback to character estimation")
+
+ def test_count_reasoning_tokens_truncated_response(self):
+ """Test counting tokens when response is truncated (no closing tag)"""
+ # Test truncated think tag
+ truncated_text = "This reasoning was cut off due to max tokens"
+
+ result1 = optillm_count(truncated_text)
+ result2 = inference_count(truncated_text)
+
+ self.assertGreater(result1, 0, "Should count tokens from truncated think block")
+ self.assertEqual(result1, result2, "Both functions should return same result")
+
+ def test_count_reasoning_tokens_mixed_complete_and_truncated(self):
+ """Test with both complete and truncated think blocks"""
+ mixed_text = """
+ First complete reasoning block
+ Some output here
+ This second block was truncated and never closed
+ """
+
+ result = optillm_count(mixed_text)
+ self.assertGreater(result, 0, "Should count tokens from both complete and truncated blocks")
+
+ # Should be more than just the first block alone
+ first_block_only = "First complete reasoning block"
+ first_result = optillm_count(first_block_only)
+ self.assertGreater(result, first_result, "Should include truncated content")
+
+ def test_count_reasoning_tokens_no_false_positives(self):
+ """Test that we don't count think-like content that isn't actually truncated"""
+ # This should NOT be counted as truncated since there's a later
+ text_with_complete_blocks = "First blockOutputSecond complete block"
+
+ result = optillm_count(text_with_complete_blocks)
+
+ # Count manually - should only be the content inside the two complete blocks
+ manual_count = optillm_count("First blockSecond complete block")
+ self.assertEqual(result, manual_count, "Should only count complete blocks, not detect false truncation")
+
+ def test_count_reasoning_tokens_edge_cases_truncated(self):
+ """Test edge cases with truncated responses"""
+ test_cases = [
+ ("", 0), # Just opening tag, no content
+ ("a", 1), # Minimal content
+ ("Some output reasoning here", None), # Truncated at end
+ ("multi\nline\ntruncated", None), # Multiline truncated
+ ]
+
+ for text, expected_min in test_cases:
+ result = optillm_count(text)
+ if expected_min is not None:
+ if expected_min == 0:
+ self.assertEqual(result, expected_min, f"Should return {expected_min} for: {text}")
+ else:
+ self.assertGreaterEqual(result, expected_min, f"Should be at least {expected_min} for: {text}")
+ else:
+ self.assertGreater(result, 0, f"Should count truncated content for: {text}")
+
+
+class TestInferenceStructures(unittest.TestCase):
+ """Test that inference structures support reasoning tokens"""
+
+ def test_chat_completion_usage_with_reasoning_tokens(self):
+ """Test ChatCompletionUsage supports reasoning_tokens"""
+ from optillm.inference import ChatCompletionUsage
+
+ # Test with reasoning tokens
+ usage = ChatCompletionUsage(
+ prompt_tokens=10,
+ completion_tokens=20,
+ total_tokens=30,
+ reasoning_tokens=5
+ )
+
+ self.assertEqual(usage.reasoning_tokens, 5)
+
+ # Test default value
+ usage_default = ChatCompletionUsage(
+ prompt_tokens=10,
+ completion_tokens=20,
+ total_tokens=30
+ )
+
+ self.assertEqual(usage_default.reasoning_tokens, 0)
+
+ def test_chat_completion_model_dump_structure(self):
+ """Test ChatCompletion model_dump includes reasoning_tokens"""
+ from optillm.inference import ChatCompletion
+
+ response_dict = {
+ "id": "test-123",
+ "object": "chat.completion",
+ "created": 1234567890,
+ "model": "test-model",
+ "choices": [{
+ "index": 0,
+ "message": {
+ "role": "assistant",
+ "content": "test response"
+ },
+ "finish_reason": "stop"
+ }],
+ "usage": {
+ "prompt_tokens": 10,
+ "completion_tokens": 15,
+ "total_tokens": 25,
+ "reasoning_tokens": 3
+ }
+ }
+
+ completion = ChatCompletion(response_dict)
+ result = completion.model_dump()
+
+ # Check structure
+ self.assertIn("usage", result)
+ self.assertIn("completion_tokens_details", result["usage"])
+ self.assertIn("reasoning_tokens", result["usage"]["completion_tokens_details"])
+ self.assertEqual(result["usage"]["completion_tokens_details"]["reasoning_tokens"], 3)
+
+
+if __name__ == '__main__':
+ # Run the tests
+ unittest.main(verbosity=2)
\ No newline at end of file
diff --git a/tests/test_reasoning_tokens.py b/tests/test_reasoning_tokens.py
new file mode 100644
index 00000000..c729d845
--- /dev/null
+++ b/tests/test_reasoning_tokens.py
@@ -0,0 +1,288 @@
+#!/usr/bin/env python3
+"""
+Tests for reasoning token functionality in OptILLM
+Covers count_reasoning_tokens function and API response format
+"""
+
+import sys
+import os
+import unittest
+from unittest.mock import Mock, patch, MagicMock
+import re
+
+# Add parent directory to path for imports
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+# Import the count_reasoning_tokens function from both modules
+from optillm import count_reasoning_tokens as optillm_count_reasoning_tokens
+from optillm.inference import count_reasoning_tokens as inference_count_reasoning_tokens
+
+
+class TestCountReasoningTokens(unittest.TestCase):
+ """Test the count_reasoning_tokens function"""
+
+ def test_count_reasoning_tokens_basic(self):
+ """Test basic functionality of count_reasoning_tokens"""
+ # Test with think tags
+ text_with_think = "This is reasoning contentThis is output"
+
+ # Test both implementations should work the same
+ result1 = optillm_count_reasoning_tokens(text_with_think)
+ result2 = inference_count_reasoning_tokens(text_with_think)
+
+ self.assertGreater(result1, 0)
+ self.assertEqual(result1, result2)
+
+ def test_count_reasoning_tokens_no_think_tags(self):
+ """Test with text that has no think tags"""
+ text_without_think = "This is just regular output text"
+
+ result1 = optillm_count_reasoning_tokens(text_without_think)
+ result2 = inference_count_reasoning_tokens(text_without_think)
+
+ self.assertEqual(result1, 0)
+ self.assertEqual(result2, 0)
+
+ def test_count_reasoning_tokens_multiple_think_blocks(self):
+ """Test with multiple think tag blocks"""
+ text_multiple = """
+ First reasoning block
+ Some output here
+ Second reasoning block with more content
+ Final output
+ """
+
+ result = optillm_count_reasoning_tokens(text_multiple)
+ self.assertGreater(result, 0)
+
+ # Should count tokens from both blocks
+ single_block = "First reasoning blockSecond reasoning block with more content"
+ single_result = optillm_count_reasoning_tokens(single_block)
+ self.assertAlmostEqual(result, single_result, delta=2) # Allow small variance due to formatting
+
+ def test_count_reasoning_tokens_empty_input(self):
+ """Test with empty or None input"""
+ self.assertEqual(optillm_count_reasoning_tokens(""), 0)
+ self.assertEqual(optillm_count_reasoning_tokens(None), 0)
+ self.assertEqual(optillm_count_reasoning_tokens(123), 0) # Non-string input
+
+ def test_count_reasoning_tokens_malformed_tags(self):
+ """Test with malformed think tags"""
+ malformed_cases = [
+ "Unclosed think tag",
+ "Unopened think tag",
+ "Nested tags",
+ "Wrong case",
+ "", # Empty think block
+ ]
+
+ for case in malformed_cases:
+ result = optillm_count_reasoning_tokens(case)
+ # Should handle gracefully, either 0 or some reasonable count
+ self.assertGreaterEqual(result, 0)
+
+ def test_count_reasoning_tokens_with_tokenizer(self):
+ """Test with a mock tokenizer for precise counting"""
+ mock_tokenizer = Mock()
+ mock_tokenizer.encode.return_value = [1, 2, 3, 4, 5] # 5 tokens
+
+ text = "Some reasoning textOutput"
+ result = optillm_count_reasoning_tokens(text, mock_tokenizer)
+
+ self.assertEqual(result, 5)
+ mock_tokenizer.encode.assert_called_once_with("Some reasoning text")
+
+ def test_count_reasoning_tokens_tokenizer_error(self):
+ """Test fallback when tokenizer fails"""
+ mock_tokenizer = Mock()
+ mock_tokenizer.encode.side_effect = Exception("Tokenizer error")
+
+ text = "Some reasoning textOutput"
+ result = optillm_count_reasoning_tokens(text, mock_tokenizer)
+
+ # Should fallback to character-based estimation
+ self.assertGreater(result, 0)
+ mock_tokenizer.encode.assert_called_once()
+
+ def test_count_reasoning_tokens_multiline(self):
+ """Test with multiline think blocks"""
+ multiline_text = """
+ This is a multi-line reasoning block
+ with several lines of content
+ that spans multiple lines
+
+ This is the final output"""
+
+ result = optillm_count_reasoning_tokens(multiline_text)
+ self.assertGreater(result, 10) # Should be substantial content
+
+ def test_count_reasoning_tokens_special_characters(self):
+ """Test with special characters in think blocks"""
+ special_text = "Content with émojis 🤔 and symbols @#$%^&*()Output"
+ result = optillm_count_reasoning_tokens(special_text)
+ self.assertGreater(result, 0)
+
+
+class TestAPIResponseFormat(unittest.TestCase):
+ """Test that API responses include reasoning token information"""
+
+ def setUp(self):
+ """Set up test fixtures"""
+ # Import after setting up path
+ import optillm
+ self.app = optillm.app
+ self.app.config['TESTING'] = True
+ self.client = self.app.test_client()
+
+ @patch('optillm.get_config')
+ def test_response_includes_completion_tokens_details(self, mock_get_config):
+ """Test that API responses include completion_tokens_details"""
+ # Mock the OpenAI client
+ mock_client = Mock()
+ mock_response = Mock()
+ mock_response.choices = [Mock()]
+ mock_response.choices[0].message.content = "Some reasoningFinal answer: 42"
+ mock_response.usage.completion_tokens = 20
+ mock_response.usage.prompt_tokens = 10
+ mock_response.usage.total_tokens = 30
+
+ mock_client.chat.completions.create.return_value = mock_response
+ mock_get_config.return_value = (mock_client, "test-key")
+
+ # Make request to the API
+ response = self.client.post('/v1/chat/completions',
+ json={
+ "model": "gpt-4o-mini",
+ "messages": [{"role": "user", "content": "What is 2+2?"}]
+ },
+ headers={"Authorization": "Bearer test-key"})
+
+ self.assertEqual(response.status_code, 200)
+
+ # Check response format
+ data = response.get_json()
+ self.assertIn('usage', data)
+ self.assertIn('completion_tokens_details', data['usage'])
+ self.assertIn('reasoning_tokens', data['usage']['completion_tokens_details'])
+ self.assertGreater(data['usage']['completion_tokens_details']['reasoning_tokens'], 0)
+
+ @patch('optillm.get_config')
+ def test_response_no_reasoning_tokens(self, mock_get_config):
+ """Test API response when there are no reasoning tokens"""
+ # Mock the OpenAI client with no think tags
+ mock_client = Mock()
+ mock_response = Mock()
+ mock_response.choices = [Mock()]
+ mock_response.choices[0].message.content = "Final answer: 42" # No think tags
+ mock_response.usage.completion_tokens = 10
+ mock_response.usage.prompt_tokens = 5
+ mock_response.usage.total_tokens = 15
+
+ mock_client.chat.completions.create.return_value = mock_response
+ mock_get_config.return_value = (mock_client, "test-key")
+
+ # Make request to the API
+ response = self.client.post('/v1/chat/completions',
+ json={
+ "model": "gpt-4o-mini",
+ "messages": [{"role": "user", "content": "What is 2+2?"}]
+ },
+ headers={"Authorization": "Bearer test-key"})
+
+ self.assertEqual(response.status_code, 200)
+
+ # Check response format
+ data = response.get_json()
+ self.assertIn('usage', data)
+ self.assertIn('completion_tokens_details', data['usage'])
+ self.assertEqual(data['usage']['completion_tokens_details']['reasoning_tokens'], 0)
+
+ @patch('optillm.get_config')
+ def test_multiple_responses_reasoning_tokens(self, mock_get_config):
+ """Test reasoning tokens with multiple responses (n > 1)"""
+ # Mock the OpenAI client with multiple responses
+ mock_client = Mock()
+ mock_response = Mock()
+
+ # Create multiple choices with different reasoning content
+ choice1 = Mock()
+ choice1.message.content = "First reasoningAnswer 1"
+ choice2 = Mock()
+ choice2.message.content = "Second longer reasoning contentAnswer 2"
+
+ mock_response.choices = [choice1, choice2]
+ mock_response.usage.completion_tokens = 30
+ mock_response.usage.prompt_tokens = 10
+ mock_response.usage.total_tokens = 40
+
+ mock_client.chat.completions.create.return_value = mock_response
+ mock_get_config.return_value = (mock_client, "test-key")
+
+ # Make request with n=2
+ response = self.client.post('/v1/chat/completions',
+ json={
+ "model": "gpt-4o-mini",
+ "messages": [{"role": "user", "content": "What is 2+2?"}],
+ "n": 2
+ },
+ headers={"Authorization": "Bearer test-key"})
+
+ self.assertEqual(response.status_code, 200)
+
+ # Check response format
+ data = response.get_json()
+ self.assertIn('usage', data)
+ self.assertIn('completion_tokens_details', data['usage'])
+ self.assertGreater(data['usage']['completion_tokens_details']['reasoning_tokens'], 0)
+
+ # Should have 2 choices
+ self.assertEqual(len(data['choices']), 2)
+
+
+class TestBackwardCompatibility(unittest.TestCase):
+ """Test backward compatibility with existing functionality"""
+
+ def test_existing_approaches_still_work(self):
+ """Test that existing approaches work without reasoning token changes"""
+ # Import approaches that don't use reasoning
+ from optillm.bon import best_of_n_sampling
+
+ # Create mock client
+ mock_client = Mock()
+ mock_response = Mock()
+ mock_response.choices = [Mock()]
+ mock_response.choices[0].message.content = "Regular response"
+ mock_response.usage.completion_tokens = 10
+
+ mock_client.chat.completions.create.return_value = mock_response
+
+ # Test that approach still works
+ try:
+ result, tokens = best_of_n_sampling(
+ model="test-model",
+ messages=[{"role": "user", "content": "test"}],
+ client=mock_client,
+ n=3
+ )
+ self.assertIsInstance(result, str)
+ self.assertIsInstance(tokens, int)
+ except Exception as e:
+ self.fail(f"Existing approach failed: {e}")
+
+ def test_api_without_auth_header(self):
+ """Test API still returns proper errors without auth"""
+ import optillm
+ app = optillm.app
+ app.config['TESTING'] = True
+ client = app.test_client()
+
+ response = client.post('/v1/chat/completions',
+ json={"model": "test", "messages": []})
+
+ # Should still return 401 for missing auth
+ self.assertEqual(response.status_code, 401)
+
+
+if __name__ == '__main__':
+ # Run the tests
+ unittest.main(verbosity=2)
\ No newline at end of file