Refactor provider generate functions to return GenerateResult

+51 -44

docs/PROVIDERS.md

··· 10 10 11 11 | Function | Purpose | 12 12 |----------|---------| 13 - | `generate()` | Synchronous text generation | 14 - | `agenerate()` | Asynchronous text generation | 13 + | `run_generate()` | Synchronous text generation, returns `GenerateResult` | 14 + | `run_agenerate()` | Asynchronous text generation, returns `GenerateResult` | 15 15 | `run_agent()` | Agentic execution with MCP tool support | 16 16 17 17 See `think/providers/__init__.py` for the canonical export list and `think/providers/google.py` as a reference implementation. ··· 51 51 52 52 **Settings app integration:** Add your provider to `PROVIDER_METADATA` in `think/providers/__init__.py` with `label` and `env_key` fields. The settings UI dynamically builds provider dropdowns from the registry. Add corresponding API key UI fields in `apps/settings/workspace.html` for user configuration. 53 53 54 - ## generate() / agenerate() 54 + ## run_generate() / run_agenerate() 55 55 56 - These functions handle direct LLM text generation. The unified API in `think/models.py` routes requests to provider-specific implementations. 56 + These functions handle direct LLM text generation. The unified API in `think/models.py` routes requests to provider-specific implementations and handles token logging and JSON validation centrally. 57 57 58 58 **Function signature:** 59 59 ```python 60 - def generate( 60 + from think.agents import GenerateResult 61 + 62 + def run_generate( 61 63 contents: Union[str, List[Any]], 62 64 model: str, 63 65 temperature: float = 0.3, ··· 66 68 json_output: bool = False, 67 69 thinking_budget: Optional[int] = None, 68 70 timeout_s: Optional[float] = None, 69 - context: Optional[str] = None, 70 71 **kwargs: Any, 71 - ) -> str: 72 + ) -> GenerateResult: 72 73 ``` 73 74 74 - The `agenerate()` function has the same signature but is `async`. 75 + The `run_agenerate()` function has the same signature but is `async`. 76 + 77 + **Return type - GenerateResult:** 78 + ```python 79 + class GenerateResult(TypedDict, total=False): 80 + text: Required[str] # Response text 81 + usage: Optional[dict] # Normalized usage dict 82 + finish_reason: Optional[str] # Normalized: "stop", "max_tokens", etc. 83 + thinking: Optional[list] # List of thinking block dicts 84 + ``` 75 85 76 86 **Parameter details:** 77 87 ··· 84 94 | `json_output` | Request JSON response. Google uses `response_mime_type`, Anthropic/OpenAI use response format or system instruction. | 85 95 | `thinking_budget` | Token budget for reasoning/thinking. Must be `> 0` to enable; `None` or `0` means no thinking. Only Google and Anthropic support this - OpenAI ignores it (uses fixed "medium" reasoning effort). Note: `run_agent()` always enables thinking regardless of this parameter. | 86 96 | `timeout_s` | Request timeout in seconds. Convert to provider's expected format (e.g., Google uses milliseconds internally). | 87 - | `context` | For token logging attribution only - not used in generation. | 88 97 | `**kwargs` | Absorb unknown kwargs for forward compatibility. Provider-specific options (e.g., `cached_content` for Google) pass through here. | 89 98 90 99 **Key responsibilities:** 91 100 - Accept the common parameter set shown above 92 - - Return the response text as a string 93 - - Call `log_token_usage()` after successful generation 94 - - Handle provider-specific response parsing and validation 101 + - Return `GenerateResult` with text, usage, finish_reason, and thinking 102 + - Normalize `finish_reason` to standard values: `"stop"`, `"max_tokens"`, `"safety"`, etc. 103 + - Handle provider-specific response parsing 104 + 105 + **Note:** Token logging and JSON validation are handled by the wrapper in `think/models.py`, not by providers. 95 106 96 107 **Important:** Providers should gracefully ignore unsupported parameters rather than raising errors. 97 108 ··· 184 195 185 196 ## Token Logging 186 197 187 - All generation calls must log token usage for cost tracking. 188 - 189 - ```python 190 - from think.models import log_token_usage 191 - 192 - log_token_usage(model=model, usage=usage_dict, context=context) 193 - ``` 198 + Token logging is handled centrally by the wrapper in `think/models.py`. Providers return usage data in their `GenerateResult`, and the wrapper calls `log_token_usage()`. 194 199 195 200 **Usage dict format:** 196 201 197 - Providers should build a usage dict from their response. The normalized format is: 202 + Providers should build a normalized usage dict from their response: 198 203 ```python 199 204 usage_dict = { 200 205 "input_tokens": 1500, # Required ··· 221 226 "total_tokens": response.usage.input_tokens + response.usage.output_tokens, 222 227 } 223 228 224 - # Google (can pass response object directly) 225 - log_token_usage(model=model, usage=response, context=context) 229 + # Google 230 + usage_dict = { 231 + "input_tokens": response.usage_metadata.prompt_token_count, 232 + "output_tokens": response.usage_metadata.candidates_token_count, 233 + "total_tokens": response.usage_metadata.total_token_count, 234 + } 226 235 ``` 227 236 228 237 **Key points:** 229 - - Call after successful generation 230 - - `log_token_usage()` normalizes various formats automatically 231 - - `context` enables attribution to specific features/operations 238 + - Return usage in `GenerateResult["usage"]` - wrapper handles logging 239 + - For `run_agent()`, include usage in the `finish` event 232 240 233 241 ## Context & Routing 234 242 ··· 302 310 303 311 ## Batch Processing 304 312 305 - The `Batch` class in `think/batch.py` automatically works with all providers via the unified `agenerate()` API. No provider-specific batch implementation is needed - just ensure your `agenerate()` works correctly. 313 + The `Batch` class in `think/batch.py` automatically works with all providers via the unified `agenerate()` API in `think/models.py`. No provider-specific batch implementation is needed - just ensure your `run_agenerate()` works correctly. 306 314 307 315 ## OpenAI-Compatible Providers 308 316 ··· 322 330 ## Checklist for New Providers 323 331 324 332 **Core implementation:** 325 - 1. Create `think/providers/<name>.py` with `__all__ = ["generate", "agenerate", "run_agent"]` 326 - 2. Implement `generate()`, `agenerate()`, `run_agent()` following signatures above 333 + 1. Create `think/providers/<name>.py` with `__all__ = ["run_generate", "run_agenerate", "run_agent"]` 334 + 2. Implement `run_generate()`, `run_agenerate()`, `run_agent()` following signatures above 335 + 3. Import `GenerateResult` from `think.agents` and return it from generate functions 327 336 328 337 **Model constants** in `think/models.py`: 329 - 3. Add model constants using the pattern `{PROVIDER}_{TIER}` (e.g., `DO_LLAMA_70B`, `DO_MISTRAL_NEMO`) 338 + 4. Add model constants using the pattern `{PROVIDER}_{TIER}` (e.g., `DO_LLAMA_70B`, `DO_MISTRAL_NEMO`) 330 339 - Existing examples: `GEMINI_FLASH`, `GPT_5`, `CLAUDE_SONNET_4` 331 - 4. Add provider tier mappings to `PROVIDER_DEFAULTS` dict 332 - 5. Update `get_model_provider()` to detect your models by prefix (critical for cost tracking) 340 + 5. Add provider tier mappings to `PROVIDER_DEFAULTS` dict 341 + 6. Update `get_model_provider()` to detect your models by prefix (critical for cost tracking) 333 342 334 - **Routing:** 335 - 6. Add `elif provider == "<name>"` cases in `think/models.py`: 336 - - `generate()` function (around line 622) 337 - - `agenerate()` function (around line 700) 338 - 7. Add routing case in `think/agents.py` `main_async()` (around line 331) 343 + **Registry:** 344 + 7. Add provider to `PROVIDER_REGISTRY` in `think/providers/__init__.py` 345 + 8. Add routing case in `think/agents.py` `main_async()` (around line 331) 339 346 340 347 **Settings UI:** 341 - 8. Add provider to `PROVIDER_METADATA` in `think/providers/__init__.py` with `label` and `env_key` 342 - 9. Add API key UI field in `apps/settings/workspace.html` 348 + 9. Add provider to `PROVIDER_METADATA` in `think/providers/__init__.py` with `label` and `env_key` 349 + 10. Add API key UI field in `apps/settings/workspace.html` 343 350 344 351 **Testing:** 345 - 10. Create unit tests in `tests/test_<name>.py` 346 - 11. Create integration tests in `tests/integration/test_<name>_backend.py` 347 - 12. Add test contexts to `fixtures/journal/config/journal.json` 352 + 11. Create unit tests in `tests/test_<name>.py` 353 + 12. Create integration tests in `tests/integration/test_<name>_backend.py` 354 + 13. Add test contexts to `fixtures/journal/config/journal.json` 348 355 349 356 **Documentation:** 350 - 13. Update `think/providers/__init__.py` docstring 351 - 14. Update `docs/THINK.md` providers table 352 - 15. Update `docs/CORTEX.md` valid provider values 357 + 14. Update `think/providers/__init__.py` docstring 358 + 15. Update `docs/THINK.md` providers table 359 + 16. Update `docs/CORTEX.md` valid provider values

+3 -3

docs/THINK.md

··· 130 130 131 131 Each provider lives in `think/providers/` and exposes a common interface: 132 132 133 - - `generate()` - Sync text generation 134 - - `agenerate()` - Async text generation 133 + - `run_generate()` - Sync text generation, returns `GenerateResult` 134 + - `run_agenerate()` - Async text generation, returns `GenerateResult` 135 135 - `run_agent()` - Agent execution with MCP tools and event streaming 136 136 137 137 For direct LLM calls, use `think.models.generate()` or `think.models.agenerate()` ··· 204 204 | Google | `think/providers/google.py` | Gemini models | 205 205 | Anthropic | `think/providers/anthropic.py` | Claude via Anthropic SDK | 206 206 207 - Providers implement `generate()`, `agenerate()`, and `run_agent()` functions. See [PROVIDERS.md](PROVIDERS.md) for implementation details. 207 + Providers implement `run_generate()`, `run_agenerate()`, and `run_agent()` functions. See [PROVIDERS.md](PROVIDERS.md) for implementation details. 208 208 209 209 ## Key Components 210 210

+17

think/agents.py

··· 113 113 ] 114 114 115 115 116 + class GenerateResult(TypedDict, total=False): 117 + """Result from provider run_generate/run_agenerate functions. 118 + 119 + Structured result that allows the wrapper to handle cross-cutting concerns 120 + like token logging and JSON validation centrally. 121 + 122 + The thinking field contains dicts with: summary (str), signature (optional str), 123 + redacted_data (optional str for Anthropic redacted thinking). 124 + """ 125 + 126 + text: Required[str] # Response text 127 + usage: Optional[dict] # Normalized usage dict (input_tokens, output_tokens, etc.) 128 + finish_reason: Optional[str] # Normalized: "stop", "max_tokens", "safety", etc. 129 + thinking: Optional[list] # List of thinking block dicts 130 + 131 + 116 132 class JSONEventWriter: 117 133 """Write JSONL events to stdout and optionally to a file.""" 118 134 ··· 261 277 "ErrorEvent", 262 278 "AgentUpdatedEvent", 263 279 "ThinkingEvent", 280 + "GenerateResult", 264 281 "Event", 265 282 "JSONEventWriter", 266 283 "JSONEventCallback",

+42 -4

think/models.py

··· 924 924 # --------------------------------------------------------------------------- 925 925 926 926 927 + def _validate_json_response(result: Dict[str, Any], json_output: bool) -> None: 928 + """Validate response for JSON output mode. 929 + 930 + Raises IncompleteJSONError if finish_reason indicates truncation. 931 + """ 932 + if not json_output: 933 + return 934 + 935 + finish_reason = result.get("finish_reason") 936 + if finish_reason and finish_reason != "stop": 937 + raise IncompleteJSONError( 938 + reason=finish_reason, 939 + partial_text=result.get("text", ""), 940 + ) 941 + 942 + 927 943 def generate( 928 944 contents: Union[str, List[Any]], 929 945 context: str, ··· 971 987 ------ 972 988 ValueError 973 989 If the resolved provider is not supported. 990 + IncompleteJSONError 991 + If json_output=True and response was truncated. 974 992 """ 975 993 from think.providers import get_provider_module 976 994 ··· 984 1002 # Get provider module via registry (raises ValueError for unknown providers) 985 1003 provider_mod = get_provider_module(provider) 986 1004 987 - return provider_mod.generate( 1005 + # Call provider's run_generate (returns GenerateResult) 1006 + result = provider_mod.run_generate( 988 1007 contents=contents, 989 1008 model=model, 990 1009 temperature=temperature, ··· 993 1012 json_output=json_output, 994 1013 thinking_budget=thinking_budget, 995 1014 timeout_s=timeout_s, 996 - context=context, 997 1015 **kwargs, 998 1016 ) 999 1017 1018 + # Validate JSON output if requested 1019 + _validate_json_response(result, json_output) 1020 + 1021 + # Log token usage centrally 1022 + if result.get("usage"): 1023 + log_token_usage(model=model, usage=result["usage"], context=context) 1024 + 1025 + return result["text"] 1026 + 1000 1027 1001 1028 async def agenerate( 1002 1029 contents: Union[str, List[Any]], ··· 1045 1072 ------ 1046 1073 ValueError 1047 1074 If the resolved provider is not supported. 1075 + IncompleteJSONError 1076 + If json_output=True and response was truncated. 1048 1077 """ 1049 1078 from think.providers import get_provider_module 1050 1079 ··· 1058 1087 # Get provider module via registry (raises ValueError for unknown providers) 1059 1088 provider_mod = get_provider_module(provider) 1060 1089 1061 - return await provider_mod.agenerate( 1090 + # Call provider's run_agenerate (returns GenerateResult) 1091 + result = await provider_mod.run_agenerate( 1062 1092 contents=contents, 1063 1093 model=model, 1064 1094 temperature=temperature, ··· 1067 1097 json_output=json_output, 1068 1098 thinking_budget=thinking_budget, 1069 1099 timeout_s=timeout_s, 1070 - context=context, 1071 1100 **kwargs, 1072 1101 ) 1102 + 1103 + # Validate JSON output if requested 1104 + _validate_json_response(result, json_output) 1105 + 1106 + # Log token usage centrally 1107 + if result.get("usage"): 1108 + log_token_usage(model=model, usage=result["usage"], context=context) 1109 + 1110 + return result["text"] 1073 1111 1074 1112 1075 1113 __all__ = [

+8 -5

think/providers/__init__.py

··· 6 6 This package contains provider-specific implementations for LLM generation 7 7 and agent execution. Each provider module exposes: 8 8 9 - - generate(): Sync text generation 10 - - agenerate(): Async text generation 9 + - run_generate(): Sync text generation, returns GenerateResult 10 + - run_agenerate(): Async text generation, returns GenerateResult 11 11 - run_agent(): Agent execution with MCP tools 12 + 13 + GenerateResult is a TypedDict with: text, usage, finish_reason, thinking. 14 + The wrapper functions in think.models handle token logging and JSON validation. 12 15 13 16 Available providers: 14 17 - google: Google Gemini models ··· 25 28 # --------------------------------------------------------------------------- 26 29 # Central registry of supported providers and their module paths. 27 30 # All registered providers must implement: 28 - # - generate(contents, model, ...) -> str 29 - # - agenerate(contents, model, ...) -> str 31 + # - run_generate(contents, model, ...) -> GenerateResult 32 + # - run_agenerate(contents, model, ...) -> GenerateResult 30 33 # - run_agent(config, on_event) -> str 31 34 # --------------------------------------------------------------------------- 32 35 ··· 61 64 Returns 62 65 ------- 63 66 ModuleType 64 - The provider module with generate, agenerate, and run_agent functions. 67 + The provider module with run_generate, run_agenerate, and run_agent functions. 65 68 66 69 Raises 67 70 ------

+75 -57

think/providers/anthropic.py

··· 5 5 """Anthropic Claude provider for agents and direct LLM generation. 6 6 7 7 This module provides the Anthropic Claude provider for the ``sol agents`` CLI 8 - and standardized generate/agenerate functions for direct LLM calls. 8 + and run_generate/run_agenerate functions returning GenerateResult. 9 9 10 10 Common Parameters 11 11 ----------------- ··· 25 25 Token budget for model thinking. 26 26 timeout_s : float, optional 27 27 Request timeout in seconds. 28 - context : str, optional 29 - Context string for token usage logging. 28 + **kwargs 29 + Additional provider-specific options. 30 30 """ 31 31 32 32 from __future__ import annotations ··· 50 50 from think.models import CLAUDE_SONNET_4 51 51 from think.utils import create_mcp_client 52 52 53 - from ..agents import JSONEventCallback, ThinkingEvent 53 + from ..agents import GenerateResult, JSONEventCallback, ThinkingEvent 54 54 55 55 # Default values are now handled internally 56 56 _DEFAULT_MODEL = CLAUDE_SONNET_4 ··· 476 476 477 477 478 478 # --------------------------------------------------------------------------- 479 - # Standardized generate/agenerate functions 479 + # run_generate / run_agenerate functions 480 480 # --------------------------------------------------------------------------- 481 481 482 482 ··· 506 506 return usage_dict 507 507 508 508 509 + def _normalize_finish_reason(stop_reason: str | None) -> str | None: 510 + """Normalize Anthropic stop_reason to standard values. 511 + 512 + Returns normalized string: "stop", "max_tokens", or None. 513 + """ 514 + if not stop_reason: 515 + return None 516 + 517 + reason = stop_reason.lower() 518 + if reason == "end_turn": 519 + return "stop" 520 + elif reason == "max_tokens": 521 + return "max_tokens" 522 + elif reason == "stop_sequence": 523 + return "stop" 524 + else: 525 + return reason 526 + 527 + 528 + def _extract_text_and_thinking(response: Any) -> tuple[str, list | None]: 529 + """Extract text and thinking blocks from Anthropic response. 530 + 531 + Returns tuple of (text, thinking_blocks). 532 + """ 533 + text = "" 534 + thinking_blocks = [] 535 + 536 + for block in response.content: 537 + if getattr(block, "type", None) == "text": 538 + text += block.text 539 + elif isinstance(block, ThinkingBlock): 540 + thinking_blocks.append( 541 + { 542 + "summary": block.thinking, 543 + "signature": block.signature, 544 + } 545 + ) 546 + elif isinstance(block, RedactedThinkingBlock): 547 + thinking_blocks.append( 548 + { 549 + "summary": "[redacted]", 550 + "redacted_data": block.data, 551 + } 552 + ) 553 + 554 + return text, thinking_blocks if thinking_blocks else None 555 + 556 + 509 557 # Cache for Anthropic clients 510 558 _anthropic_client = None 511 559 _async_anthropic_client = None ··· 552 600 return [{"role": "user", "content": str(contents)}] 553 601 554 602 555 - def generate( 603 + def run_generate( 556 604 contents: Any, 557 605 model: str = _DEFAULT_MODEL, 558 606 temperature: float = 0.3, ··· 561 609 json_output: bool = False, 562 610 thinking_budget: int | None = None, 563 611 timeout_s: float | None = None, 564 - context: str | None = None, 565 612 **kwargs: Any, 566 - ) -> str: 613 + ) -> GenerateResult: 567 614 """Generate text synchronously. 568 615 616 + Returns GenerateResult with text, usage, finish_reason, and thinking. 569 617 See module docstring for parameter details. 570 618 """ 571 - from think.models import log_token_usage 572 - 573 619 client = _get_anthropic_client() 574 620 messages = _convert_contents_to_messages(contents) 575 621 ··· 602 648 if timeout_s: 603 649 request_kwargs["timeout"] = timeout_s 604 650 605 - from think.models import IncompleteJSONError 606 - 607 651 response = client.messages.create(**request_kwargs) 608 652 609 - # Extract text first (may be partial if truncated) 610 - text = "" 611 - for block in response.content: 612 - if getattr(block, "type", None) == "text": 613 - text += block.text 653 + text, thinking = _extract_text_and_thinking(response) 654 + return GenerateResult( 655 + text=text, 656 + usage=_extract_usage_dict(response), 657 + finish_reason=_normalize_finish_reason(response.stop_reason), 658 + thinking=thinking, 659 + ) 614 660 615 - # Validate stop reason for JSON output 616 - if json_output and response.stop_reason != "end_turn": 617 - raise IncompleteJSONError( 618 - reason=response.stop_reason or "unknown", partial_text=text 619 - ) 620 661 621 - # Log token usage 622 - usage_dict = _extract_usage_dict(response) 623 - if usage_dict: 624 - log_token_usage(model=model, usage=usage_dict, context=context) 625 - 626 - return text 627 - 628 - 629 - async def agenerate( 662 + async def run_agenerate( 630 663 contents: Any, 631 664 model: str = _DEFAULT_MODEL, 632 665 temperature: float = 0.3, ··· 635 668 json_output: bool = False, 636 669 thinking_budget: int | None = None, 637 670 timeout_s: float | None = None, 638 - context: str | None = None, 639 671 **kwargs: Any, 640 - ) -> str: 672 + ) -> GenerateResult: 641 673 """Generate text asynchronously. 642 674 675 + Returns GenerateResult with text, usage, finish_reason, and thinking. 643 676 See module docstring for parameter details. 644 677 """ 645 - from think.models import log_token_usage 646 - 647 678 client = _get_async_anthropic_client() 648 679 messages = _convert_contents_to_messages(contents) 649 680 ··· 676 707 if timeout_s: 677 708 request_kwargs["timeout"] = timeout_s 678 709 679 - from think.models import IncompleteJSONError 680 - 681 710 response = await client.messages.create(**request_kwargs) 682 711 683 - # Extract text first (may be partial if truncated) 684 - text = "" 685 - for block in response.content: 686 - if getattr(block, "type", None) == "text": 687 - text += block.text 688 - 689 - # Validate stop reason for JSON output 690 - if json_output and response.stop_reason != "end_turn": 691 - raise IncompleteJSONError( 692 - reason=response.stop_reason or "unknown", partial_text=text 693 - ) 694 - 695 - # Log token usage 696 - usage_dict = _extract_usage_dict(response) 697 - if usage_dict: 698 - log_token_usage(model=model, usage=usage_dict, context=context) 699 - 700 - return text 712 + text, thinking = _extract_text_and_thinking(response) 713 + return GenerateResult( 714 + text=text, 715 + usage=_extract_usage_dict(response), 716 + finish_reason=_normalize_finish_reason(response.stop_reason), 717 + thinking=thinking, 718 + ) 701 719 702 720 703 721 __all__ = [ 704 722 "run_agent", 705 - "generate", 706 - "agenerate", 723 + "run_generate", 724 + "run_agenerate", 707 725 ]

+93 -43

think/providers/google.py

··· 5 5 """Gemini provider for agents and direct LLM generation. 6 6 7 7 This module provides the Google Gemini provider for the ``sol agents`` CLI 8 - and standardized generate/agenerate functions for direct LLM calls. 8 + and run_generate/run_agenerate functions returning GenerateResult. 9 9 10 10 Common Parameters 11 11 ----------------- ··· 25 25 Token budget for model thinking. 26 26 timeout_s : float, optional 27 27 Request timeout in seconds. 28 - context : str, optional 29 - Context string for token usage logging. 30 28 **kwargs 31 29 Provider-specific options (cached_content, client). 32 30 """ ··· 46 44 from think.models import GEMINI_FLASH 47 45 from think.utils import create_mcp_client 48 46 49 - from ..agents import JSONEventCallback, ThinkingEvent 47 + from ..agents import GenerateResult, JSONEventCallback, ThinkingEvent 50 48 51 49 _DEFAULT_MAX_TOKENS = 8192 52 50 _DEFAULT_MODEL = GEMINI_FLASH ··· 145 143 return types.GenerateContentConfig(**config_args) 146 144 147 145 148 - def _validate_response(response: Any, json_output: bool = False) -> str: 149 - """Validate response and extract text. 146 + def _extract_response_text(response: Any) -> str: 147 + """Extract text from response. 150 148 151 - Returns response.text if available, or a user-friendly message for empty 152 - responses (e.g., tool-only completions). Raises on actual errors. 149 + Returns response.text if available, or a friendly completion message 150 + if the response is empty. Raises on safety filter blocks. 153 151 154 152 Parameters 155 153 ---------- 156 154 response 157 155 The response from the model. 158 - json_output 159 - If True, validates that finish_reason is STOP to ensure complete JSON. 160 156 """ 161 157 if response is None: 162 158 raise ValueError("No response from model") 163 159 164 - from think.models import IncompleteJSONError 165 - 166 160 # Check for error conditions in candidates 167 161 finish_reason = _extract_finish_reason(response) 168 162 if finish_reason and "SAFETY" in finish_reason.upper(): 169 163 raise ValueError(f"Response blocked by safety filters: {finish_reason}") 170 164 171 - # Extract text (may be partial if truncated) 165 + # Extract text, or generate friendly message if empty 172 166 text = response.text if response.text else "" 173 - 174 - # For JSON output, require STOP to ensure complete response 175 - if json_output: 176 - normalized = (finish_reason or "").upper().replace("FINISHREASON.", "") 177 - if normalized != "STOP": 178 - raise IncompleteJSONError( 179 - reason=finish_reason or "unknown", partial_text=text 180 - ) 181 - 182 - # Return text or user-friendly completion message 183 167 if text: 184 168 return text 185 169 186 - # Empty text - generate appropriate message (no tools in generate/agenerate) 170 + # Empty text - generate user-friendly completion message 187 171 return _format_completion_message(finish_reason, had_tool_calls=False) 188 172 189 173 174 + def _normalize_finish_reason(response: Any) -> str | None: 175 + """Normalize finish_reason to standard values. 176 + 177 + Returns normalized string: "stop", "max_tokens", "safety", or None. 178 + """ 179 + raw = _extract_finish_reason(response) 180 + if not raw: 181 + return None 182 + 183 + # Normalize (handle both enum names and string values) 184 + reason = raw.upper().replace("FINISHREASON.", "") 185 + 186 + if reason == "STOP": 187 + return "stop" 188 + elif reason == "MAX_TOKENS": 189 + return "max_tokens" 190 + elif "SAFETY" in reason: 191 + return "safety" 192 + elif reason == "RECITATION": 193 + return "recitation" 194 + else: 195 + return reason.lower() 196 + 197 + 198 + def _extract_usage(response: Any) -> dict | None: 199 + """Extract normalized usage dict from response.""" 200 + if not hasattr(response, "usage_metadata") or not response.usage_metadata: 201 + return None 202 + 203 + metadata = response.usage_metadata 204 + usage: dict[str, int] = { 205 + "input_tokens": getattr(metadata, "prompt_token_count", 0), 206 + "output_tokens": getattr(metadata, "candidates_token_count", 0), 207 + "total_tokens": getattr(metadata, "total_token_count", 0), 208 + } 209 + # Only include optional fields if non-zero 210 + cached = getattr(metadata, "cached_content_token_count", 0) 211 + if cached: 212 + usage["cached_tokens"] = cached 213 + reasoning = getattr(metadata, "thoughts_token_count", 0) 214 + if reasoning: 215 + usage["reasoning_tokens"] = reasoning 216 + return usage 217 + 218 + 219 + def _extract_thinking(response: Any) -> list | None: 220 + """Extract thinking blocks from response. 221 + 222 + Returns list of ThinkingBlock dicts or None if no thinking. 223 + """ 224 + if not hasattr(response, "candidates") or not response.candidates: 225 + return None 226 + 227 + thinking_blocks = [] 228 + for candidate in response.candidates: 229 + if not candidate.content or not candidate.content.parts: 230 + continue 231 + for part in candidate.content.parts: 232 + if getattr(part, "thought", False) and getattr(part, "text", None): 233 + thinking_blocks.append({"summary": part.text}) 234 + 235 + return thinking_blocks if thinking_blocks else None 236 + 237 + 190 238 def _extract_finish_reason(response: Any) -> str | None: 191 239 """Extract finish_reason from response candidates. 192 240 ··· 290 338 291 339 292 340 # --------------------------------------------------------------------------- 293 - # Standardized generate/agenerate functions 341 + # run_generate / run_agenerate functions 294 342 # --------------------------------------------------------------------------- 295 343 296 344 297 - def generate( 345 + def run_generate( 298 346 contents: str | list[Any], 299 347 model: str = _DEFAULT_MODEL, 300 348 temperature: float = 0.3, ··· 303 351 json_output: bool = False, 304 352 thinking_budget: int | None = None, 305 353 timeout_s: float | None = None, 306 - context: str | None = None, 307 354 **kwargs: Any, 308 - ) -> str: 355 + ) -> GenerateResult: 309 356 """Generate text synchronously. 310 357 358 + Returns GenerateResult with text, usage, finish_reason, and thinking. 311 359 See module docstring for parameter details. 312 360 """ 313 - from think.models import log_token_usage 314 - 315 361 cached_content = kwargs.get("cached_content") 316 362 client = kwargs.get("client") 317 363 ··· 334 380 config=config, 335 381 ) 336 382 337 - text = _validate_response(response, json_output=json_output) 338 - log_token_usage(model=model, usage=response, context=context) 339 - return text 383 + return GenerateResult( 384 + text=_extract_response_text(response), 385 + usage=_extract_usage(response), 386 + finish_reason=_normalize_finish_reason(response), 387 + thinking=_extract_thinking(response), 388 + ) 340 389 341 390 342 - async def agenerate( 391 + async def run_agenerate( 343 392 contents: str | list[Any], 344 393 model: str = _DEFAULT_MODEL, 345 394 temperature: float = 0.3, ··· 348 397 json_output: bool = False, 349 398 thinking_budget: int | None = None, 350 399 timeout_s: float | None = None, 351 - context: str | None = None, 352 400 **kwargs: Any, 353 - ) -> str: 401 + ) -> GenerateResult: 354 402 """Generate text asynchronously. 355 403 404 + Returns GenerateResult with text, usage, finish_reason, and thinking. 356 405 See module docstring for parameter details. 357 406 """ 358 - from think.models import log_token_usage 359 - 360 407 cached_content = kwargs.get("cached_content") 361 408 client = kwargs.get("client") 362 409 ··· 379 426 config=config, 380 427 ) 381 428 382 - text = _validate_response(response, json_output=json_output) 383 - log_token_usage(model=model, usage=response, context=context) 384 - return text 429 + return GenerateResult( 430 + text=_extract_response_text(response), 431 + usage=_extract_usage(response), 432 + finish_reason=_normalize_finish_reason(response), 433 + thinking=_extract_thinking(response), 434 + ) 385 435 386 436 387 437 # --------------------------------------------------------------------------- ··· 751 801 752 802 __all__ = [ 753 803 "run_agent", 754 - "generate", 755 - "agenerate", 804 + "run_generate", 805 + "run_agenerate", 756 806 "get_or_create_client", 757 807 ]

+56 -59

think/providers/openai.py

··· 5 5 """OpenAI provider for agents and direct LLM generation. 6 6 7 7 This module provides the OpenAI provider for the ``sol agents`` CLI 8 - and standardized generate/agenerate functions for direct LLM calls. 8 + and run_generate/run_agenerate functions returning GenerateResult. 9 9 10 10 Common Parameters 11 11 ----------------- ··· 21 21 Whether to request JSON response format. 22 22 timeout_s : float, optional 23 23 Request timeout in seconds. 24 - context : str, optional 25 - Context string for token usage logging. 24 + **kwargs 25 + Additional provider-specific options. 26 26 27 27 Note: GPT-5+ reasoning models don't support custom temperature (fixed at 1.0) 28 28 and reasoning is always enabled with medium effort. ··· 69 69 70 70 from think.models import GPT_5 71 71 72 - from ..agents import JSONEventCallback, ThinkingEvent 72 + from ..agents import GenerateResult, JSONEventCallback, ThinkingEvent 73 73 74 74 75 75 def _convert_turns_to_items(turns: list[dict]) -> list[dict]: ··· 518 518 519 519 520 520 # --------------------------------------------------------------------------- 521 - # Standardized generate/agenerate functions 521 + # run_generate / run_agenerate functions 522 522 # --------------------------------------------------------------------------- 523 523 524 524 # Cache for OpenAI clients ··· 581 581 return messages 582 582 583 583 584 - def generate( 584 + def _normalize_finish_reason(finish_reason: str | None) -> str | None: 585 + """Normalize OpenAI finish_reason to standard values. 586 + 587 + Returns normalized string: "stop", "max_tokens", "content_filter", or None. 588 + """ 589 + if not finish_reason: 590 + return None 591 + 592 + reason = finish_reason.lower() 593 + if reason == "stop": 594 + return "stop" 595 + elif reason == "length": 596 + return "max_tokens" 597 + elif reason == "content_filter": 598 + return "content_filter" 599 + else: 600 + return reason 601 + 602 + 603 + def _extract_usage(response: Any) -> dict | None: 604 + """Extract normalized usage dict from OpenAI response.""" 605 + if not response.usage: 606 + return None 607 + 608 + return { 609 + "input_tokens": response.usage.prompt_tokens, 610 + "output_tokens": response.usage.completion_tokens, 611 + "total_tokens": response.usage.total_tokens, 612 + } 613 + 614 + 615 + def run_generate( 585 616 contents: Any, 586 617 model: str = GPT_5, 587 618 max_output_tokens: int = 8192 * 2, 588 619 system_instruction: str | None = None, 589 620 json_output: bool = False, 590 621 timeout_s: float | None = None, 591 - context: str | None = None, 592 622 **kwargs: Any, 593 - ) -> str: 623 + ) -> GenerateResult: 594 624 """Generate text synchronously. 595 625 626 + Returns GenerateResult with text, usage, finish_reason, and thinking. 596 627 See module docstring for parameter details. 597 628 """ 598 - from think.models import log_token_usage 599 - 600 629 client = _get_openai_client() 601 630 messages = _convert_contents_to_messages(contents, system_instruction) 602 631 ··· 616 645 if timeout_s: 617 646 request_kwargs["timeout"] = timeout_s 618 647 619 - from think.models import IncompleteJSONError 620 - 621 648 response = client.chat.completions.create(**request_kwargs) 622 649 623 - # Extract text first (may be partial if truncated) 624 650 choice = response.choices[0] 625 - text = choice.message.content or "" 651 + return GenerateResult( 652 + text=choice.message.content or "", 653 + usage=_extract_usage(response), 654 + finish_reason=_normalize_finish_reason(choice.finish_reason), 655 + thinking=None, # OpenAI reasoning not exposed in generate response 656 + ) 626 657 627 - # Validate finish reason for JSON output 628 - if json_output and choice.finish_reason != "stop": 629 - raise IncompleteJSONError( 630 - reason=choice.finish_reason or "unknown", partial_text=text 631 - ) 632 658 633 - # Log token usage 634 - if response.usage: 635 - usage_dict = { 636 - "input_tokens": response.usage.prompt_tokens, 637 - "output_tokens": response.usage.completion_tokens, 638 - "total_tokens": response.usage.total_tokens, 639 - } 640 - log_token_usage(model=model, usage=usage_dict, context=context) 641 - 642 - return text 643 - 644 - 645 - async def agenerate( 659 + async def run_agenerate( 646 660 contents: Any, 647 661 model: str = GPT_5, 648 662 max_output_tokens: int = 8192 * 2, 649 663 system_instruction: str | None = None, 650 664 json_output: bool = False, 651 665 timeout_s: float | None = None, 652 - context: str | None = None, 653 666 **kwargs: Any, 654 - ) -> str: 667 + ) -> GenerateResult: 655 668 """Generate text asynchronously. 656 669 670 + Returns GenerateResult with text, usage, finish_reason, and thinking. 657 671 See module docstring for parameter details. 658 672 """ 659 - from think.models import log_token_usage 660 - 661 673 client = _get_async_openai_client() 662 674 messages = _convert_contents_to_messages(contents, system_instruction) 663 675 ··· 677 689 if timeout_s: 678 690 request_kwargs["timeout"] = timeout_s 679 691 680 - from think.models import IncompleteJSONError 681 - 682 692 response = await client.chat.completions.create(**request_kwargs) 683 693 684 - # Extract text first (may be partial if truncated) 685 694 choice = response.choices[0] 686 - text = choice.message.content or "" 687 - 688 - # Validate finish reason for JSON output 689 - if json_output and choice.finish_reason != "stop": 690 - raise IncompleteJSONError( 691 - reason=choice.finish_reason or "unknown", partial_text=text 692 - ) 693 - 694 - # Log token usage 695 - if response.usage: 696 - usage_dict = { 697 - "input_tokens": response.usage.prompt_tokens, 698 - "output_tokens": response.usage.completion_tokens, 699 - "total_tokens": response.usage.total_tokens, 700 - } 701 - log_token_usage(model=model, usage=usage_dict, context=context) 702 - 703 - return text 695 + return GenerateResult( 696 + text=choice.message.content or "", 697 + usage=_extract_usage(response), 698 + finish_reason=_normalize_finish_reason(choice.finish_reason), 699 + thinking=None, # OpenAI reasoning not exposed in generate response 700 + ) 704 701 705 702 706 703 __all__ = [ 707 704 "run_agent", 708 - "generate", 709 - "agenerate", 705 + "run_generate", 706 + "run_agenerate", 710 707 ]

Configure Feed

Configure Feed