excpose reasoning tokesn count

Author: codelion

optillm/__init__.py

@@ -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',

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()

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__))))

tests/test_plugins.py

@@ -79,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"