Merge pull request #341 from enoch3712/327-bug-litellm-max_tokens-no-supported

[BUG] Litellm: max_tokens no supported

Author: enoch3712

extract_thinker/concatenation_handler.py

@@ -13,6 +13,9 @@ def __init__(self, llm):
 
     def _is_valid_json_continuation(self, response: str) -> bool:
         """Check if the response is a valid JSON continuation."""
+        if not response:
+            return False
+        
         cleaned_response = response.strip()
 
         # Check if response contains JSON markers

extract_thinker/extractor.py

@@ -215,6 +215,9 @@ def extract(
         Returns:
             The parsed result from the LLM as validated by response_model.
         """
+        if isinstance(source, dict) and self.document_loader is None:
+            self.document_loader = DocumentLoaderData()
+        
         self._validate_dependencies(response_model, vision)
         self.extra_content = content
         self.completion_strategy = completion_strategy

extract_thinker/global_models.py

@@ -1,9 +1,21 @@
 def get_lite_model():
     """Return the lite model for cost efficiency."""
     #return "vertex_ai/gemini-2.0-flash"
-    return "gemini/gemini-2.0-flash"
+    return "gemini/gemini-2.5-flash-preview-05-20"
 
 def get_big_model():
     """Return the big model for high performance."""
     #return "vertex_ai/gemini-2.0-flash"
-    return "gemini/gemini-2.0-flash"
+    return "gemini/gemini-2.5-flash-preview-05-20"
+
+def get_gemini_flash_model():
+    """Return the Gemini 2.5 Flash model."""
+    return "gemini/gemini-2.5-flash-preview-05-20"
+
+def get_gpt_mini_model():
+    """Return the GPT-4.1 Mini model."""
+    return "gpt-4.1-mini"
+
+def get_gpt_o4_model():
+    """Return the GPT-4o model."""
+    return "gpt-4o"

extract_thinker/llm.py

@@ -32,6 +32,7 @@ class LLM:
     MAX_TOKEN_LIMIT = 120000  # Maximum token limit (for Claude 3.7 Sonnet)
     MAX_THINKING_BUDGET = 64000  # Maximum thinking budget
     MIN_THINKING_BUDGET = 1200  # Minimum thinking budget
+    DEFAULT_OUTPUT_TOKENS = 32000
 
     def __init__(
         self,
@@ -55,6 +56,7 @@ def __init__(
         self.is_thinking = False  # Initialize is_thinking flag
         self.page_count = None  # Initialize page count
         self.thinking_budget = self.THINKING_BUDGET_TOKENS  # Default thinking budget
+        self.thinking_token_limit: Optional[int] = None
 
         if self.backend == LLMEngine.DEFAULT:
             self.client = instructor.from_litellm(
@@ -67,7 +69,6 @@ def __init__(
             from pydantic_ai import Agent
             from pydantic_ai.models import KnownModelName
             from typing import cast
-            import asyncio
 
             self.client = None
             self.agent = Agent(
@@ -158,7 +159,7 @@ def set_page_count(self, page_count: int) -> None:
         thinking_tokens = min(thinking_tokens, self.MAX_THINKING_BUDGET)
 
         # Update token limit and thinking budget
-        self.token_limit = content_tokens
+        self.thinking_token_limit = content_tokens
         self.thinking_budget = thinking_tokens
 
     def request(
@@ -171,14 +172,7 @@ def request(
             # Combine messages into a single prompt
             combined_prompt = " ".join([m["content"] for m in messages])
             try:
-                # Create event loop if it doesn't exist
-                try:
-                    loop = asyncio.get_event_loop()
-                except RuntimeError:
-                    loop = asyncio.new_event_loop()
-                    asyncio.set_event_loop(loop)
-                
-                result = loop.run_until_complete(
+                result = asyncio.run(
                     self.agent.run(
                         combined_prompt, 
                         result_type=response_model if response_model else str
@@ -225,139 +219,108 @@ def request(
 
     def _request_with_router(self, messages: List[Dict[str, str]], response_model: Optional[str]) -> Any:
         """Handle request using router with or without thinking parameter"""
+        max_tokens = self.DEFAULT_OUTPUT_TOKENS
+        if self.token_limit is not None:
+            max_tokens = self.token_limit
+        elif self.is_thinking:
+            max_tokens = self.thinking_token_limit
+        
+        params = {
+            "model": self.model,
+            "messages": messages,
+            "response_model": response_model,
+            "temperature": self.temperature,
+            "timeout": self.TIMEOUT,
+            "max_completion_tokens": max_tokens,
+        }
         if self.is_thinking:
-            # Add thinking parameter for supported models
-            thinking_param = {
-                "type": "enabled",
-                "budget_tokens": self.thinking_budget
-            }
-            try:
-                return self.router.completion(
-                    model=self.model,
-                    messages=messages,
-                    response_model=response_model,
-                    temperature=self.temperature,
-                    timeout=self.TIMEOUT,
-                    thinking=thinking_param,
-                )
-            except Exception as e:
-                # If thinking parameter causes an error, try without it
-                if "property 'thinking' is unsupported" in str(e):
-                    print(f"Warning: Model {self.model} doesn't support thinking parameter, proceeding without it.")
-                    return self.router.completion(
-                        model=self.model,
-                        messages=messages,
-                        response_model=response_model,
-                        temperature=self.temperature,
-                        timeout=self.TIMEOUT,
-                    )
-                else:
-                    raise e
-        else:
-            # Normal request without thinking parameter
-            return self.router.completion(
-                model=self.model,
-                messages=messages,
-                response_model=response_model,
-                temperature=self.temperature,
-                timeout=self.TIMEOUT,
-            )
+            if litellm.supports_reasoning(self.model):
+                # Add thinking parameter for supported models
+                thinking_param = {
+                    "type": "enabled",
+                    "budget_tokens": self.thinking_budget
+                }
+                params["thinking"] = thinking_param
+            else:
+                print(f"Warning: Model {self.model} doesn't support thinking parameter, proceeding without it.")
+
+        return self.router.completion(**params)
 
     def _request_direct(self, messages: List[Dict[str, str]], response_model: Optional[str]) -> Any:
         """Handle direct request with or without thinking parameter"""
+        max_tokens = self.DEFAULT_OUTPUT_TOKENS
+        if self.token_limit is not None:
+            max_tokens = self.token_limit
+        elif self.is_thinking:
+            max_tokens = self.thinking_token_limit
+
         base_params = {
             "model": self.model,
             "messages": messages,
             "temperature": self.temperature,
             "response_model": response_model,
             "max_retries": 1,
-            "max_tokens": self.token_limit,
+            "max_completion_tokens": max_tokens,
             "timeout": self.TIMEOUT,
         }
 
         if self.is_thinking:
-            # Try with thinking parameter
-            thinking_param = {
-                "type": "enabled",
-                "budget_tokens": self.thinking_budget
-            }
-            try:
-                return self.client.chat.completions.create(
-                    **base_params,
-                    thinking=thinking_param,
-                )
-            except Exception as e:
-                # If thinking parameter causes an error, try without it
-                if "property 'thinking' is unsupported" in str(e):
-                    print(f"Warning: Model {self.model} doesn't support thinking parameter, proceeding without it.")
-                    return self.client.chat.completions.create(**base_params)
-                else:
-                    raise e
-        else:
-            # Normal request without thinking parameter
-            return self.client.chat.completions.create(**base_params)
-
-    def raw_completion(self, messages: List[Dict[str, str]]) -> str:
-        """Make raw completion request without response model."""
-        if self.router:
-            if self.is_thinking:
-                # Add thinking parameter for supported models
+            if litellm.supports_reasoning(self.model):
+                # Try with thinking parameter
                 thinking_param = {
                     "type": "enabled",
                     "budget_tokens": self.thinking_budget
                 }
-                try:
-                    raw_response = self.router.completion(
-                        model=self.model,
-                        messages=messages,
-                        thinking=thinking_param,
-                    )
-                except Exception as e:
-                    # If thinking parameter causes an error, try without it
-                    if "property 'thinking' is unsupported" in str(e):
-                        print(f"Warning: Model {self.model} doesn't support thinking parameter, proceeding without it.")
-                        raw_response = self.router.completion(
-                            model=self.model,
-                            messages=messages,
-                        )
-                    else:
-                        raise e
+                base_params["thinking"] = thinking_param
             else:
-                raw_response = self.router.completion(
-                    model=self.model,
-                    messages=messages,
+                print(f"Warning: Model {self.model} doesn't support thinking parameter, proceeding without it.")
+        
+        return self.client.chat.completions.create(**base_params)
+
+    def raw_completion(self, messages: List[Dict[str, str]]) -> str:
+        """Make raw completion request without response model."""
+        if self.backend == LLMEngine.PYDANTIC_AI:
+            # Combine messages into a single prompt
+            combined_prompt = " ".join([m["content"] for m in messages])
+            try:
+                result = asyncio.run(
+                    self.agent.run(
+                        combined_prompt, 
+                        result_type=str
+                    )
                 )
-        else:
-            if self.is_thinking:
+                return result.data
+            except Exception as e:
+                raise ValueError(f"Failed to extract from source: {str(e)}")
+
+        max_tokens = self.DEFAULT_OUTPUT_TOKENS
+        if self.token_limit is not None:
+            max_tokens = self.token_limit
+        elif self.is_thinking:
+            max_tokens = self.thinking_token_limit
+
+        params = {
+            "model": self.model,
+            "messages": messages,
+            "max_completion_tokens": max_tokens,
+        }
+
+        if self.is_thinking:
+            if litellm.supports_reasoning(self.model):
                 # Add thinking parameter for supported models
                 thinking_param = {
                     "type": "enabled",
                     "budget_tokens": self.thinking_budget
                 }
-                try:
-                    raw_response = litellm.completion(
-                        model=self.model,
-                        messages=messages,
-                        max_tokens=self.token_limit,
-                        thinking=thinking_param,
-                    )
-                except Exception as e:
-                    # If thinking parameter causes an error, try without it
-                    if "property 'thinking' is unsupported" in str(e):
-                        print(f"Warning: Model {self.model} doesn't support thinking parameter, proceeding without it.")
-                        raw_response = litellm.completion(
-                            model=self.model,
-                            messages=messages,
-                            max_tokens=self.token_limit,
-                        )
-                    else:
-                        raise e
+                params["thinking"] = thinking_param
             else:
-                raw_response = litellm.completion(
-                    model=self.model,
-                    messages=messages,
-                    max_tokens=self.token_limit,
-                )
+                print(f"Warning: Model {self.model} doesn't support thinking parameter, proceeding without it.")
+        
+        if self.router:
+            raw_response = self.router.completion(**params)
+        else:
+            raw_response = litellm.completion(**params)
+        
         return raw_response.choices[0].message.content
 
     def set_timeout(self, timeout_ms: int) -> None: