From 111d25cd59023f79c4165c199c075cf848b43e1a Mon Sep 17 00:00:00 2001 From: Yinghuan Zhang Date: Sun, 14 Jun 2026 13:42:09 -0400 Subject: [PATCH 1/8] Rename model-client chat() to response() and add solve_physics_problem() Align the low-level model-client call with OpenAI's responses.create: rename chat() -> response(), rename the prompt parameter user_prompt -> input, and add a keyword-only `instructions` (system prompt) parameter. When `instructions` is omitted every provider except OpenAI falls back to a short DEFAULT_INSTRUCTIONS system prompt (resolved via the new _resolve_instructions helper, which OpenAI overrides to stay input-only); an explicit empty string suppresses it. Each provider injects the system prompt at its native site (Anthropic top-level `system`, OpenAI `instructions`, openai-compatible/Ollama system message, Gemini `system_instruction`). chat_structured threads `instructions` through; a deprecated chat() alias warns and forwards to response(). Add a high-level BaseModelClient.solve_physics_problem() that dispatches on the input type: a plain str question, a PhysicsProblem (formatted into a prompt plus its image_path images), or a PhysicsQuestionSemantics (NotImplementedError TODO, lazy-imported to keep core off semantics at import time); any other type raises TypeError. Output is gated by the new PhysicsOutputMode enum (only ANSWER_TEXT implemented). Add core-only modes.py and prompts.py; the semantics _format_problem now reuses format_problem_context for its shared header. Centralize the annotation worker's physics-expert system prompt via `instructions` and update the calls.py retry path. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/prkit/annotation/workers/base.py | 32 +-- src/prkit/core/model_clients/__init__.py | 8 +- src/prkit/core/model_clients/anthropic.py | 15 +- src/prkit/core/model_clients/base.py | 139 +++++++++++- src/prkit/core/model_clients/dashscope.py | 11 +- src/prkit/core/model_clients/deepseek.py | 11 +- src/prkit/core/model_clients/gemini.py | 15 +- src/prkit/core/model_clients/modes.py | 26 +++ src/prkit/core/model_clients/ollama.py | 20 +- src/prkit/core/model_clients/openai.py | 24 +- .../model_clients/openai_compatible_chat.py | 29 ++- src/prkit/core/model_clients/prompts.py | 49 ++++ src/prkit/semantics/inference/calls.py | 4 +- src/prkit/semantics/inference/prompts.py | 33 +-- tests/prkit/annotation/workers/test_base.py | 12 +- .../annotation/workers/test_domain_labeler.py | 8 +- .../workers/test_theorem_detector.py | 8 +- .../workers/test_variable_locator.py | 8 +- .../core/model_clients/test_anthropic.py | 36 ++- tests/prkit/core/model_clients/test_base.py | 211 +++++++++++++++--- .../core/model_clients/test_dashscope.py | 23 +- .../prkit/core/model_clients/test_deepseek.py | 22 +- tests/prkit/core/model_clients/test_gemini.py | 36 ++- tests/prkit/core/model_clients/test_ollama.py | 40 ++-- tests/prkit/core/model_clients/test_openai.py | 44 +++- .../test_openai_compatible_chat.py | 8 +- .../core/model_clients/test_package_init.py | 15 ++ .../prkit/core/model_clients/test_prompts.py | 64 ++++++ tests/prkit/core/model_clients/test_xai.py | 6 +- .../prkit/semantics/test_inference_prompts.py | 24 +- 30 files changed, 779 insertions(+), 202 deletions(-) create mode 100644 src/prkit/core/model_clients/modes.py create mode 100644 src/prkit/core/model_clients/prompts.py create mode 100644 tests/prkit/core/model_clients/test_prompts.py diff --git a/src/prkit/annotation/workers/base.py b/src/prkit/annotation/workers/base.py index 35dece3..47d0a69 100644 --- a/src/prkit/annotation/workers/base.py +++ b/src/prkit/annotation/workers/base.py @@ -3,7 +3,7 @@ from abc import ABC, abstractmethod from typing import Any -from prkit.core.model_clients import create_model_client +from prkit.core.model_clients import DEFAULT_INSTRUCTIONS, create_model_client from prkit.core.model_clients.base import BaseModelClient @@ -51,19 +51,17 @@ def _call_llm_structured(self, prompt: str, response_format: Any) -> Any: or None if call fails """ try: - full_prompt = ( - "You are a physics expert. Provide accurate, detailed analysis of physics problems. " - "Always respond with valid JSON in the exact format requested.\n\n" - f"{prompt}" - ) if hasattr(response_format, "model_validate") and ( "chat_structured" in type(self.llm_client).__dict__ or isinstance(self.llm_client, BaseModelClient) ): + # The physics-expert role is sent as `instructions`; the JSON + # contract is added by chat_structured's structured-output suffix. result = self.llm_client.chat_structured( - full_prompt, + prompt, response_model=response_format, structured_policy="best_effort", + instructions=DEFAULT_INSTRUCTIONS, ) if result.parsed is not None: return result.parsed @@ -74,7 +72,14 @@ def _call_llm_structured(self, prompt: str, response_format: Any) -> Any: return response_format(**response_dict) return None - response_text = self.llm_client.chat(full_prompt) + # Fallback for clients without structured output: ask for JSON inline. + json_prompt = ( + f"{prompt}\n\n" + "Always respond with valid JSON in the exact format requested." + ) + response_text = self.llm_client.response( + input=json_prompt, instructions=DEFAULT_INSTRUCTIONS + ) if response_text: import json @@ -98,12 +103,13 @@ def _call_llm(self, prompt: str) -> str: Response text from LLM, or empty JSON string if call fails """ try: - full_prompt = ( - "You are a physics expert. Provide accurate, detailed analysis of physics problems. " - "Always respond with valid JSON in the exact format requested.\n\n" - f"{prompt}" + json_prompt = ( + f"{prompt}\n\n" + "Always respond with valid JSON in the exact format requested." ) - return self.llm_client.chat(full_prompt).strip() + return self.llm_client.response( + input=json_prompt, instructions=DEFAULT_INSTRUCTIONS + ).strip() except Exception as e: print(f"Error calling LLM API: {e}") return "{}" diff --git a/src/prkit/core/model_clients/__init__.py b/src/prkit/core/model_clients/__init__.py index 4680cb7..a3f93c4 100644 --- a/src/prkit/core/model_clients/__init__.py +++ b/src/prkit/core/model_clients/__init__.py @@ -8,16 +8,22 @@ from typing import Any -from .base import BaseModelClient +from .base import DEFAULT_INSTRUCTIONS, BaseModelClient from .factory import ProviderRule, create_model_client, register_model_client +from .modes import PhysicsOutputMode +from .prompts import build_plain_question_prompt, format_problem_context create_llm_client = create_model_client __all__ = [ "BaseModelClient", + "DEFAULT_INSTRUCTIONS", + "PhysicsOutputMode", "ProviderRule", + "build_plain_question_prompt", "create_model_client", "create_llm_client", + "format_problem_context", "register_model_client", "AnthropicModel", "DashscopeModel", diff --git a/src/prkit/core/model_clients/anthropic.py b/src/prkit/core/model_clients/anthropic.py index 2ecc7b7..f90715d 100644 --- a/src/prkit/core/model_clients/anthropic.py +++ b/src/prkit/core/model_clients/anthropic.py @@ -183,20 +183,21 @@ def __init__(self, model: str, logger: logging.Logger | None = None) -> None: self.client: Any = Anthropic(api_key=os.environ.get("ANTHROPIC_API_KEY")) self.provider = "anthropic" - def chat( + def response( self, - user_prompt: str, + input: str, image_paths: list[str] | None = None, response_format: dict[str, Any] | type | None = None, max_output_tokens: int = 1024, *args: Any, + instructions: str | None = None, **kwargs: Any, ) -> str: """ Generate a response from Anthropic Messages API. Args: - user_prompt: The user's prompt text (string) + input: The user's prompt text (string) image_paths: Optional list of image paths/URLs (strings). Supports: - File paths: encoded to base64 - Base64 data URLs: passed as-is after parsing @@ -205,6 +206,9 @@ def chat( request is translated into a forced Anthropic tool call. max_output_tokens: Maximum output tokens for Anthropic API. *args: Additional positional arguments (ignored, kept for compatibility) + instructions: Optional system prompt. Sent as the top-level Anthropic + ``system`` parameter. Defaults to ``DEFAULT_INSTRUCTIONS`` + when omitted (see ``_resolve_instructions``). **kwargs: Additional keyword arguments for request parameters (e.g., temperature, top_p, etc.) @@ -219,7 +223,7 @@ def chat( if response_format is not None: normalized_response_format = normalize_response_format(response_format) - content: list[dict[str, Any]] = [{"type": "text", "text": user_prompt}] + content: list[dict[str, Any]] = [{"type": "text", "text": input}] if image_paths: for image_path in image_paths: if image_path.startswith("data:"): @@ -261,6 +265,9 @@ def chat( "messages": [{"role": "user", "content": content}], "max_tokens": max_output_tokens, } + instr = self._resolve_instructions(instructions) + if instr: + request_params["system"] = instr if normalized_response_format is not None: request_params["output_config"] = { "format": { diff --git a/src/prkit/core/model_clients/base.py b/src/prkit/core/model_clients/base.py index 1119ddc..1f29191 100644 --- a/src/prkit/core/model_clients/base.py +++ b/src/prkit/core/model_clients/base.py @@ -3,14 +3,18 @@ from __future__ import annotations import logging +import warnings from abc import ABC, abstractmethod from collections.abc import Sequence -from typing import Any, TypeVar +from typing import TYPE_CHECKING, Any, TypeVar from pydantic import BaseModel, ValidationError +from ..domain import PhysicsProblem from ..logging_config import PRKitLogger from ..project_env import load_project_dotenv +from .modes import PhysicsOutputMode +from .prompts import build_plain_question_prompt from .structured_output import ( StructuredCallResult, StructuredOutputPlan, @@ -22,8 +26,19 @@ normalize_response_format, ) +if TYPE_CHECKING: + from prkit.semantics.schema import PhysicsQuestionSemantics + T = TypeVar("T", bound=BaseModel) +# Default system prompt sent to every provider except OpenAI when the caller +# does not supply their own ``instructions``. Kept short and free of any +# response-format directives — structured output is handled separately via the +# per-provider ``response_format`` / prompt-suffix machinery. +DEFAULT_INSTRUCTIONS = ( + "You are a physics expert. Provide accurate, detailed analysis of physics problems." +) + class BaseModelClient(ABC): """Abstract base class for all model client implementations.""" @@ -46,6 +61,16 @@ def _provider_name(self) -> str: """Return the provider identifier string, or ``'unknown'`` when unset.""" return self.provider or "unknown" + def _resolve_instructions(self, instructions: str | None) -> str | None: + """Resolve the effective system prompt (``instructions``) to send. + + Returns *instructions* when the caller supplied one (including an empty + string, which suppresses the system prompt), otherwise falls back to + :data:`DEFAULT_INSTRUCTIONS`. Providers that do not require a system + prompt (e.g. OpenAI) override this to return *instructions* unchanged. + """ + return instructions if instructions is not None else DEFAULT_INSTRUCTIONS + @property def supports_native_structured_output(self) -> bool: """True when this provider supports at least one form of native structured output.""" @@ -55,6 +80,24 @@ def supports_native_structured_output(self) -> bool: ) @abstractmethod + def response( + self, + input: str, + image_paths: list[str] | None = None, + response_format: dict[str, Any] | type | None = None, + *, + instructions: str | None = None, + **kwargs: Any, + ) -> str: + """Send a single request to the provider and return the response text. + + Mirrors OpenAI's ``client.responses.create``: *input* is the user prompt + and *instructions* is the (optional) system prompt. When *instructions* + is ``None``, non-OpenAI providers fall back to + :data:`DEFAULT_INSTRUCTIONS` (see :meth:`_resolve_instructions`). + """ + raise NotImplementedError("Subclasses must implement .response()") + def chat( self, user_prompt: str, @@ -62,7 +105,19 @@ def chat( response_format: dict[str, Any] | type | None = None, **kwargs: Any, ) -> str: - raise NotImplementedError("Subclasses must implement .chat()") + """Deprecated alias for :meth:`response`; use ``response()`` instead.""" + warnings.warn( + "BaseModelClient.chat() is deprecated; use .response() instead " + "(the first parameter is now `input` rather than `user_prompt`).", + DeprecationWarning, + stacklevel=2, + ) + return self.response( + input=user_prompt, + image_paths=image_paths, + response_format=response_format, + **kwargs, + ) def resolve_structured_output_plan( self, @@ -98,6 +153,7 @@ def chat_structured( image_paths: Sequence[str] | None = None, max_output_tokens: int | None = None, structured_policy: StructuredOutputPolicy = "best_effort", + instructions: str | None = None, **kwargs: Any, ) -> StructuredCallResult[T]: """Send a chat request and parse the response into a Pydantic model instance.""" @@ -110,10 +166,11 @@ def chat_structured( if max_output_tokens is not None: request_kwargs["max_output_tokens"] = max_output_tokens - raw_text = self.chat( - user_prompt=prompt, + raw_text = self.response( + input=prompt, image_paths=list(image_paths) if image_paths else None, response_format=plan.response_format, + instructions=instructions, **request_kwargs, ) return self._build_structured_call_result( @@ -122,6 +179,80 @@ def chat_structured( plan=plan, ) + def solve_physics_problem( + self, + problem: PhysicsProblem | PhysicsQuestionSemantics | str, + *, + output_mode: PhysicsOutputMode = PhysicsOutputMode.ANSWER_TEXT, + instructions: str | None = None, + response_format: dict[str, Any] | type | None = None, + **kwargs: Any, + ) -> str: + """Ask the model to solve a physics problem and return its answer. + + The *input kind* is dispatched on the runtime type of *problem*: + + - ``str``: treated directly as the question text (no images). + - :class:`~prkit.core.domain.PhysicsProblem`: parsed into a prompt and its + attached ``image_path`` images. + - :class:`~prkit.semantics.schema.PhysicsQuestionSemantics`: physics + question semantics input — not yet supported (raises + ``NotImplementedError``). + - any other type: unsupported (raises ``TypeError``). + + *output_mode* selects what the model returns. Only + :attr:`PhysicsOutputMode.ANSWER_TEXT` is implemented; answer-semantics + output raises ``NotImplementedError``. The return type is ``str`` today + and will broaden to a structured result once answer semantics lands. + + This builds the prompt and prepares images, then delegates to + :meth:`response`. + """ + if output_mode is not PhysicsOutputMode.ANSWER_TEXT: + # TODO(answer-semantics-output): build the prediction-answer-semantics + # request via prkit.semantics and return a structured result. + raise NotImplementedError( + f"output_mode={output_mode!r} is not implemented yet; only " + f"{PhysicsOutputMode.ANSWER_TEXT!r} is supported." + ) + + if isinstance(problem, str): + prompt = problem.strip() + image_paths: list[str] | None = None + elif isinstance(problem, PhysicsProblem): + prompt = build_plain_question_prompt(problem) + image_paths = problem.image_path or None + else: + # Question-semantics input is a TODO; everything else is unsupported. + try: + from prkit.semantics.schema import PhysicsQuestionSemantics + except ImportError: + question_semantics_types: tuple[type, ...] = () + else: + question_semantics_types = (PhysicsQuestionSemantics,) + + if question_semantics_types and isinstance( + problem, question_semantics_types + ): + # TODO(question-semantics-input): render question semantics into a + # prompt via prkit.semantics. + raise NotImplementedError( + "Question-semantics input is not supported yet; pass a " + "PhysicsProblem or a plain question string." + ) + raise TypeError( + f"Unsupported problem input type: {type(problem)!r}. Expected a " + "PhysicsProblem, a question string, or a PhysicsQuestionSemantics." + ) + + return self.response( + input=prompt, + image_paths=image_paths, + response_format=response_format, + instructions=instructions, + **kwargs, + ) + def build_batch_structured_request( self, *, diff --git a/src/prkit/core/model_clients/dashscope.py b/src/prkit/core/model_clients/dashscope.py index df7316c..62bc768 100644 --- a/src/prkit/core/model_clients/dashscope.py +++ b/src/prkit/core/model_clients/dashscope.py @@ -113,11 +113,13 @@ def _default_enable_thinking(self) -> bool | None: return None - def chat( + def response( self, - user_prompt: str, + input: str, image_paths: list[str] | None = None, response_format: dict[str, Any] | type | None = None, + *, + instructions: str | None = None, **kwargs: Any, ) -> str: """Send a request to DashScope, automatically managing the thinking-mode flag.""" @@ -147,10 +149,11 @@ def chat( if extra_body: kwargs["extra_body"] = extra_body - return super().chat( - user_prompt=user_prompt, + return super().response( + input=input, image_paths=image_paths, response_format=response_format, + instructions=instructions, **kwargs, ) diff --git a/src/prkit/core/model_clients/deepseek.py b/src/prkit/core/model_clients/deepseek.py index 058155a..3a06d32 100644 --- a/src/prkit/core/model_clients/deepseek.py +++ b/src/prkit/core/model_clients/deepseek.py @@ -103,17 +103,20 @@ def _resolve_structured_output_plan( ), ) - def chat( + def response( self, - user_prompt: str, + input: str, image_paths: list[str] | None = None, response_format: dict[str, Any] | type | None = None, + *, + instructions: str | None = None, **kwargs: Any, ) -> str: """Send a chat request to DeepSeek and return the model's text response.""" - return super().chat( - user_prompt=user_prompt, + return super().response( + input=input, image_paths=image_paths, response_format=response_format, + instructions=instructions, **kwargs, ) diff --git a/src/prkit/core/model_clients/gemini.py b/src/prkit/core/model_clients/gemini.py index 0e98b2c..6487efe 100644 --- a/src/prkit/core/model_clients/gemini.py +++ b/src/prkit/core/model_clients/gemini.py @@ -43,26 +43,30 @@ def __init__(self, model: str, logger: logging.Logger | None = None) -> None: self.genai_client = genai.Client() self.provider = "google" - def chat( + def response( self, - user_prompt: str, + input: str, image_paths: list[str] | None = None, response_format: dict[str, Any] | type | None = None, max_output_tokens: int = 65535, *args: Any, + instructions: str | None = None, **kwargs: Any, ) -> str: """ Generate a response from Gemini API. Args: - user_prompt: The user's prompt text (string) + input: The user's prompt text (string) image_paths: Optional list of image paths/URLs (strings). Note: Gemini models support vision, but this implementation currently only handles text. Images are ignored with a warning. response_format: Optional structured output format (OpenAI-style dict or Pydantic model). Converted to Gemini's response_json_schema. *args: Additional positional arguments (ignored, kept for compatibility) + instructions: Optional system prompt, sent as the Gemini + ``system_instruction`` config field. Defaults to + ``DEFAULT_INSTRUCTIONS`` when omitted. **kwargs: Additional keyword arguments for generate_content config (e.g., temperature, max_tokens, etc.) @@ -71,7 +75,7 @@ def chat( """ # Prepare content parts # Start with the text prompt - contents_parts: list[Any] = [user_prompt] + contents_parts: list[Any] = [input] # Process images if provided if image_paths: @@ -90,6 +94,9 @@ def chat( # Build config with any additional kwargs config_dict: dict[str, Any] = {"max_output_tokens": max_output_tokens} + instr = self._resolve_instructions(instructions) + if instr: + config_dict["system_instruction"] = instr if kwargs: config_dict.update(kwargs) diff --git a/src/prkit/core/model_clients/modes.py b/src/prkit/core/model_clients/modes.py new file mode 100644 index 0000000..db5694b --- /dev/null +++ b/src/prkit/core/model_clients/modes.py @@ -0,0 +1,26 @@ +"""Mode enums for the high-level ``solve_physics_problem`` client method.""" + +from __future__ import annotations + +from enum import Enum + + +class _StrEnum(str, Enum): + """Enum subclass with string values and friendly ``str()`` output.""" + + def __str__(self) -> str: + return str(self.value) + + +class PhysicsOutputMode(_StrEnum): + """What form the model's answer should take. + + The input *kind* is determined by the runtime type of the argument passed to + ``solve_physics_problem`` (plain ``str``, ``PhysicsProblem``, or — once + supported — a physics ``PhysicsQuestionSemantics``), so there is no + corresponding input-mode enum. + """ + + ANSWER_TEXT = "answer_text" + # TODO: structured physics answer semantics output (needs prkit.semantics). + ANSWER_SEMANTICS = "answer_semantics" diff --git a/src/prkit/core/model_clients/ollama.py b/src/prkit/core/model_clients/ollama.py index c603bad..ad2aa90 100644 --- a/src/prkit/core/model_clients/ollama.py +++ b/src/prkit/core/model_clients/ollama.py @@ -156,20 +156,21 @@ def _check_ollama_running(self) -> None: self.logger.error(error_msg) raise ConnectionError(error_msg) from e - def chat( + def response( self, - user_prompt: str, + input: str, image_paths: list[str] | None = None, response_format: dict[str, Any] | type | None = None, max_output_tokens: int = 65535, *args: Any, + instructions: str | None = None, **kwargs: Any, ) -> str: """ Generate a response using the local Ollama service. Args: - user_prompt: The user's prompt text. + input: The user's prompt text. image_paths: Optional list of file paths to images. Ollama accepts local file paths. response_format: Optional structured output format. Supports either: @@ -177,6 +178,8 @@ def chat( - A Pydantic BaseModel class - {"type": "json_object"} for generic JSON mode *args: Additional positional arguments (ignored, kept for compatibility) + instructions: Optional system prompt, prepended as a ``system`` message. + Defaults to ``DEFAULT_INSTRUCTIONS`` when omitted. **kwargs: Additional keyword arguments for request parameters (e.g., max_tokens, etc.) @@ -190,7 +193,7 @@ def chat( """ message: dict[str, Any] = { "role": "user", - "content": user_prompt, + "content": input, } request_format: str | dict[str, Any] | None = None @@ -215,6 +218,13 @@ def chat( valid_images.append(path) message["images"] = valid_images + + messages: list[dict[str, Any]] = [] + instr = self._resolve_instructions(instructions) + if instr: + messages.append({"role": "system", "content": instr}) + messages.append(message) + options: dict[str, Any] = { "temperature": 0, "num_predict": max_output_tokens, @@ -226,7 +236,7 @@ def chat( try: request_kwargs: dict[str, Any] = { "model": self.model, - "messages": [message], + "messages": messages, "options": options, } if request_format is not None: diff --git a/src/prkit/core/model_clients/openai.py b/src/prkit/core/model_clients/openai.py index 5d42739..cf926af 100644 --- a/src/prkit/core/model_clients/openai.py +++ b/src/prkit/core/model_clients/openai.py @@ -232,19 +232,28 @@ def __init__( self.base_url = base_url self.is_o_family = _is_o_family_model(model) - def chat( + def _resolve_instructions(self, instructions: str | None) -> str | None: + """OpenAI works with ``input`` alone — no default system prompt. + + Only the caller's explicit *instructions* (if any) are sent; ``None`` stays + ``None`` so the Responses API receives no ``instructions`` parameter. + """ + return instructions + + def response( self, - user_prompt: str, + input: str, image_paths: list[str] | None = None, response_format: dict[str, Any] | type | None = None, *args: Any, + instructions: str | None = None, **kwargs: Any, ) -> str: """ Generate a response from OpenAI Responses API. Args: - user_prompt: The user's prompt text (string) + input: The user's prompt text (string) image_paths: Optional list of image paths/URLs (strings). Can be: - File paths: ["/path/to/image.jpg", ...] - will be encoded to base64 - HTTP/HTTPS URLs: ["https://example.com/image.jpg", ...] - used as-is @@ -252,6 +261,9 @@ def chat( response_format: Optional structured output format (OpenAI-style dict or Pydantic model). Ensures response adheres to JSON Schema. *args: Additional positional arguments (ignored, kept for compatibility) + instructions: Optional system prompt, sent as the Responses API + ``instructions`` parameter. OpenAI does not apply a default, + so it is omitted entirely when not provided. **kwargs: Additional keyword arguments for request parameters (e.g., max_tokens, etc.) @@ -280,7 +292,7 @@ def chat( text_format["description"] = normalized["description"] # Use role/content format for all models - content: list[dict[str, Any]] = [{"type": "input_text", "text": user_prompt}] + content: list[dict[str, Any]] = [{"type": "input_text", "text": input}] if image_paths: for image_path in image_paths: @@ -289,6 +301,10 @@ def chat( request_params["input"] = [{"role": "user", "content": content}] + instr = self._resolve_instructions(instructions) + if instr: + request_params["instructions"] = instr + # Add reasoning parameter for o-family models if self.is_o_family: request_params["reasoning"] = {"effort": "medium"} diff --git a/src/prkit/core/model_clients/openai_compatible_chat.py b/src/prkit/core/model_clients/openai_compatible_chat.py index a8776a5..2c8e57a 100644 --- a/src/prkit/core/model_clients/openai_compatible_chat.py +++ b/src/prkit/core/model_clients/openai_compatible_chat.py @@ -140,25 +140,32 @@ def _extract_text_from_chat_completion(self, response: Any) -> str: return "\n".join(chunk for chunk in text_chunks if chunk).strip() return str(content) - def chat( + def response( self, - user_prompt: str, + input: str, image_paths: list[str] | None = None, response_format: dict[str, Any] | type | None = None, + *, + instructions: str | None = None, **kwargs: Any, ) -> str: """Send a chat-completions request and return the model's text response.""" + messages: list[dict[str, Any]] = [] + instr = self._resolve_instructions(instructions) + if instr: + messages.append({"role": "system", "content": instr}) + messages.append( + { + "role": "user", + "content": self._build_message_content( + self._structured_prompt_for_chat(input, response_format), + image_paths, + ), + } + ) request_params: dict[str, Any] = { "model": self.model, - "messages": [ - { - "role": "user", - "content": self._build_message_content( - self._structured_prompt_for_chat(user_prompt, response_format), - image_paths, - ), - } - ], + "messages": messages, } request_response_format = self._build_chat_response_format(response_format) diff --git a/src/prkit/core/model_clients/prompts.py b/src/prkit/core/model_clients/prompts.py new file mode 100644 index 0000000..7a522c8 --- /dev/null +++ b/src/prkit/core/model_clients/prompts.py @@ -0,0 +1,49 @@ +"""Core-only prompt builders for asking a model to solve a physics problem. + +These helpers depend only on :mod:`prkit.core.domain` so the model-client layer +stays self-contained (no dependency on the higher-level ``prkit.semantics`` +package). The semantics package reuses :func:`format_problem_context` for the +shared problem-context header. +""" + +from __future__ import annotations + +from ..domain import PhysicsProblem + + +def format_problem_context(problem: PhysicsProblem) -> str: + """Render the shared problem-context header embedded in solver prompts. + + Includes the problem id, type, domain, language, an attached-images notice, + multiple-choice options, and the question text. It deliberately excludes any + reference answer or worked solution so it is safe to send to a model that is + being asked to solve the problem. + """ + lines = [ + f"Problem ID: {problem.problem_id}", + f"Problem type: {problem.problem_type or 'unspecified'}", + f"Domain: {problem.get_domain_name()}", + f"Language: {problem.language}", + ] + + if problem.image_path: + lines.append( + f"Attached images: {len(problem.image_path)} image(s) are provided " + "separately with this request." + ) + + if problem.options: + option_lines = [] + for index, option in enumerate(problem.options): + label = chr(ord("A") + index) + option_lines.append(f"{label}. {option}") + lines.append("Options:\n" + "\n".join(option_lines)) + + lines.append("Question:\n" + (problem.question or "").strip()) + + return "\n".join(lines) + + +def build_plain_question_prompt(problem: PhysicsProblem) -> str: + """Build the prompt for solving *problem* from its plain question text.""" + return format_problem_context(problem) diff --git a/src/prkit/semantics/inference/calls.py b/src/prkit/semantics/inference/calls.py index 74f3f20..54e9c91 100644 --- a/src/prkit/semantics/inference/calls.py +++ b/src/prkit/semantics/inference/calls.py @@ -1043,8 +1043,8 @@ def _retry_non_native_json_completion( getattr(model_client, "model", "unknown"), getattr(model_client, "provider", "unknown"), ) - return model_client.chat( - user_prompt=retry_prompt, + return model_client.response( + input=retry_prompt, image_paths=list(image_paths) or None, response_format=None, **retry_kwargs, diff --git a/src/prkit/semantics/inference/prompts.py b/src/prkit/semantics/inference/prompts.py index ca94d4d..09271f4 100644 --- a/src/prkit/semantics/inference/prompts.py +++ b/src/prkit/semantics/inference/prompts.py @@ -3,6 +3,7 @@ from __future__ import annotations from prkit.core.domain import Answer, PhysicsProblem +from prkit.core.model_clients.prompts import format_problem_context from ..normalization import ( infer_prediction_question_semantics, @@ -142,39 +143,25 @@ def _format_problem( *, include_reference_context: bool = True, ) -> str: - """Render the problem context block that is embedded in prompts.""" + """Render the problem context block that is embedded in prompts. - lines = [ - f"Problem ID: {problem.problem_id}", - f"Problem type: {problem.problem_type or 'unspecified'}", - f"Domain: {problem.get_domain_name()}", - f"Language: {problem.language}", - ] - - if problem.image_path: - lines.append( - f"Attached images: {len(problem.image_path)} image(s) are provided separately with this request." - ) - - if problem.options: - option_lines = [] - for index, option in enumerate(problem.options): - label = chr(ord("A") + index) - option_lines.append(f"{label}. {option}") - lines.append("Options:\n" + "\n".join(option_lines)) + The shared header (id/type/domain/language/images/options/question) is + produced by the core-layer ``format_problem_context``; the reference-context + (answer/solution) block is appended here only when requested. + """ - lines.append("Question:\n" + (problem.question or "").strip()) + sections = [format_problem_context(problem)] if include_reference_context: answer_text = answer_like_to_text(problem.answer) if answer_text: - lines.append("Answer:\n" + answer_text) + sections.append("Answer:\n" + answer_text) solution_text = _problem_solution_text(problem) if solution_text: - lines.append("Solution:\n" + solution_text) + sections.append("Solution:\n" + solution_text) - return "\n".join(lines) + return "\n".join(sections) def _problem_solution_text(problem: PhysicsProblem) -> str: diff --git a/tests/prkit/annotation/workers/test_base.py b/tests/prkit/annotation/workers/test_base.py index a1806e8..0394d61 100644 --- a/tests/prkit/annotation/workers/test_base.py +++ b/tests/prkit/annotation/workers/test_base.py @@ -52,7 +52,7 @@ class TestResponse(BaseModel): result: str mock_client = Mock() - mock_client.chat.return_value = '{"result": "structured"}' + mock_client.response.return_value = '{"result": "structured"}' mock_create.return_value = mock_client annotator = ConcreteAnnotator(model="gpt-5.1") @@ -60,13 +60,13 @@ class TestResponse(BaseModel): assert result is not None assert result.result == "structured" - mock_client.chat.assert_called_once() + mock_client.response.assert_called_once() @patch("prkit.annotation.workers.base.create_model_client") def test_base_annotator_call_llm_structured_error(self, mock_create): """Test _call_llm_structured with error handling.""" mock_client = Mock() - mock_client.chat.side_effect = Exception("API error") + mock_client.response.side_effect = Exception("API error") mock_create.return_value = mock_client annotator = ConcreteAnnotator(model="gpt-4o") @@ -78,20 +78,20 @@ def test_base_annotator_call_llm_structured_error(self, mock_create): def test_base_annotator_call_llm(self, mock_create): """Test _call_llm method.""" mock_client = Mock() - mock_client.chat.return_value = "LLM response" + mock_client.response.return_value = "LLM response" mock_create.return_value = mock_client annotator = ConcreteAnnotator(model="gpt-4o") result = annotator._call_llm("test prompt") assert result == "LLM response" - mock_client.chat.assert_called_once() + mock_client.response.assert_called_once() @patch("prkit.annotation.workers.base.create_model_client") def test_base_annotator_call_llm_error(self, mock_create): """Test _call_llm with error handling.""" mock_client = Mock() - mock_client.chat.side_effect = Exception("API error") + mock_client.response.side_effect = Exception("API error") mock_create.return_value = mock_client annotator = ConcreteAnnotator(model="gpt-4o") diff --git a/tests/prkit/annotation/workers/test_domain_labeler.py b/tests/prkit/annotation/workers/test_domain_labeler.py index 984a55b..7026283 100644 --- a/tests/prkit/annotation/workers/test_domain_labeler.py +++ b/tests/prkit/annotation/workers/test_domain_labeler.py @@ -36,7 +36,7 @@ def test_domain_labeler_work_success(self, mock_create): import json # Return JSON string that will be parsed by _call_llm_structured - mock_client.chat.return_value = json.dumps(mock_response.model_dump()) + mock_client.response.return_value = json.dumps(mock_response.model_dump()) mock_create.return_value = mock_client labeler = DomainLabeler(model="gpt-5.1") @@ -59,7 +59,7 @@ def test_domain_labeler_work_with_invalid_domain(self, mock_create): import json # Return JSON string that will be parsed by _call_llm_structured - mock_client.chat.return_value = json.dumps(mock_response.model_dump()) + mock_client.response.return_value = json.dumps(mock_response.model_dump()) mock_create.return_value = mock_client labeler = DomainLabeler(model="gpt-5.1") @@ -73,7 +73,7 @@ def test_domain_labeler_work_with_invalid_domain(self, mock_create): def test_domain_labeler_work_llm_error(self, mock_create): """Test domain labeling when LLM call fails.""" mock_client = Mock() - mock_client.chat.return_value = None # Simulate failure + mock_client.response.return_value = None # Simulate failure mock_create.return_value = mock_client labeler = DomainLabeler(model="gpt-5.1") @@ -98,7 +98,7 @@ def test_domain_labeler_normalizes_domain_strings(self, mock_create): import json # Return JSON string that will be parsed by _call_llm_structured - mock_client.chat.return_value = json.dumps(mock_response.model_dump()) + mock_client.response.return_value = json.dumps(mock_response.model_dump()) mock_create.return_value = mock_client labeler = DomainLabeler(model="gpt-5.1") diff --git a/tests/prkit/annotation/workers/test_theorem_detector.py b/tests/prkit/annotation/workers/test_theorem_detector.py index 0a9dddf..4a29d5e 100644 --- a/tests/prkit/annotation/workers/test_theorem_detector.py +++ b/tests/prkit/annotation/workers/test_theorem_detector.py @@ -42,7 +42,7 @@ def test_theorem_detector_work_success(self, mock_create): # _call_llm_structured now parses JSON and returns Pydantic model import json - mock_client.chat.return_value = json.dumps(mock_response.model_dump()) + mock_client.response.return_value = json.dumps(mock_response.model_dump()) mock_create.return_value = mock_client detector = TheoremDetector(model="gpt-5.1") @@ -58,7 +58,7 @@ def test_theorem_detector_work_fallback(self, mock_create): """Test theorem detection falls back to regular LLM call.""" mock_client = Mock() # Structured call fails (now just chat with combined prompt) - mock_client.chat.side_effect = [ + mock_client.response.side_effect = [ Exception("Structured failed"), # First call fails json.dumps( # Fallback call succeeds { @@ -85,7 +85,7 @@ def test_theorem_detector_work_fallback(self, mock_create): def test_theorem_detector_work_both_fail(self, mock_create): """Test theorem detection when both methods fail.""" mock_client = Mock() - mock_client.chat.side_effect = Exception("All calls failed") + mock_client.response.side_effect = Exception("All calls failed") mock_create.return_value = mock_client detector = TheoremDetector(model="gpt-5.1") @@ -99,7 +99,7 @@ def test_theorem_detector_work_both_fail(self, mock_create): def test_theorem_detector_work_invalid_json(self, mock_create): """Test theorem detection with invalid JSON in fallback.""" mock_client = Mock() - mock_client.chat.side_effect = [ + mock_client.response.side_effect = [ None, # First call returns None (treated as failure) "Invalid JSON {", # Fallback call returns invalid JSON ] diff --git a/tests/prkit/annotation/workers/test_variable_locator.py b/tests/prkit/annotation/workers/test_variable_locator.py index 8ac6bc0..3545348 100644 --- a/tests/prkit/annotation/workers/test_variable_locator.py +++ b/tests/prkit/annotation/workers/test_variable_locator.py @@ -52,7 +52,7 @@ def test_variable_locator_work_success(self, mock_create): mock_response = VariableResponse( variables=variables, problem_summary="Find time given velocity" ) - mock_client.chat.return_value = json.dumps(mock_response.model_dump()) + mock_client.response.return_value = json.dumps(mock_response.model_dump()) mock_create.return_value = mock_client locator = VariableLocator(model="gpt-5.1") @@ -69,7 +69,7 @@ def test_variable_locator_work_success(self, mock_create): def test_variable_locator_work_fallback(self, mock_create): """Test variable extraction falls back to regular LLM call.""" mock_client = Mock() - mock_client.chat.side_effect = [ + mock_client.response.side_effect = [ Exception("Structured failed"), # First call fails json.dumps( # Fallback call succeeds { @@ -97,7 +97,7 @@ def test_variable_locator_work_fallback(self, mock_create): def test_variable_locator_work_both_fail(self, mock_create): """Test variable extraction when both methods fail.""" mock_client = Mock() - mock_client.chat.side_effect = Exception("All calls failed") + mock_client.response.side_effect = Exception("All calls failed") mock_create.return_value = mock_client locator = VariableLocator(model="gpt-5.1") @@ -121,7 +121,7 @@ def test_variable_locator_separates_known_unknown(self, mock_create): ), ] mock_response = VariableResponse(variables=variables, problem_summary="Test") - mock_client.chat.return_value = json.dumps(mock_response.model_dump()) + mock_client.response.return_value = json.dumps(mock_response.model_dump()) mock_create.return_value = mock_client locator = VariableLocator(model="gpt-5.1") diff --git a/tests/prkit/core/model_clients/test_anthropic.py b/tests/prkit/core/model_clients/test_anthropic.py index 425c47e..2e5f341 100644 --- a/tests/prkit/core/model_clients/test_anthropic.py +++ b/tests/prkit/core/model_clients/test_anthropic.py @@ -13,6 +13,7 @@ _extract_tool_use_json, _parse_data_url, ) +from prkit.core.model_clients.base import DEFAULT_INSTRUCTIONS ANTHROPIC_TEST_MODEL = "claude-sonnet-4-6" @@ -48,13 +49,15 @@ def test_chat_text_only(self, mock_anthropic_class): mock_client.messages.create.return_value = mock_response client = AnthropicModel(ANTHROPIC_TEST_MODEL) - response = client.chat("Hello, world!") + response = client.response("Hello, world!") assert response == "Test response" mock_client.messages.create.assert_called_once() call_kwargs = mock_client.messages.create.call_args[1] assert call_kwargs["model"] == ANTHROPIC_TEST_MODEL assert call_kwargs["max_tokens"] == 1024 + # Default system prompt is sent as Anthropic's top-level `system` param. + assert call_kwargs["system"] == DEFAULT_INSTRUCTIONS assert call_kwargs["messages"] == [ { "role": "user", @@ -62,6 +65,27 @@ def test_chat_text_only(self, mock_anthropic_class): } ] + @patch("prkit.core.model_clients.anthropic.Anthropic") + def test_response_system_prompt_explicit_and_suppressed(self, mock_anthropic_class): + """Explicit instructions override the default; an empty string suppresses it.""" + mock_client = MagicMock() + mock_anthropic_class.return_value = mock_client + + text_block = Mock() + text_block.type = "text" + text_block.text = "ok" + mock_response = Mock() + mock_response.content = [text_block] + mock_client.messages.create.return_value = mock_response + + client = AnthropicModel(ANTHROPIC_TEST_MODEL) + + client.response("Hi", instructions="Custom system.") + assert mock_client.messages.create.call_args[1]["system"] == "Custom system." + + client.response("Hi", instructions="") + assert "system" not in mock_client.messages.create.call_args[1] + @patch("prkit.core.model_clients.anthropic.Anthropic") def test_chat_with_data_url_image(self, mock_anthropic_class): """Test chat with base64 data URL image input.""" @@ -77,7 +101,7 @@ def test_chat_with_data_url_image(self, mock_anthropic_class): client = AnthropicModel(ANTHROPIC_TEST_MODEL) data_url = "data:image/png;base64,ZmFrZS1kYXRh" - response = client.chat("Describe image", image_paths=[data_url]) + response = client.response("Describe image", image_paths=[data_url]) assert response == "Image response" call_kwargs = mock_client.messages.create.call_args[1] @@ -104,7 +128,7 @@ def test_chat_with_http_image_warns_and_ignores(self, mock_anthropic_class): client = AnthropicModel(ANTHROPIC_TEST_MODEL) with patch.object(client.logger, "warning") as mock_warning: - response = client.chat( + response = client.response( "Hello", image_paths=["https://example.com/image.jpg"], ) @@ -133,7 +157,7 @@ class ExampleResponse(BaseModel): mock_client.messages.create.return_value = mock_response client = AnthropicModel(ANTHROPIC_TEST_MODEL) - response = client.chat( + response = client.response( "Hello", response_format=ExampleResponse, ) @@ -157,7 +181,7 @@ def test_chat_invalid_data_url_raises_value_error(self, mock_anthropic_class): client = AnthropicModel(ANTHROPIC_TEST_MODEL) with pytest.raises(ValueError, match="base64"): - client.chat("Hello", image_paths=["data:image/png,not-base64"]) + client.response("Hello", image_paths=["data:image/png,not-base64"]) @patch("prkit.core.model_clients.anthropic.encode_image_to_base64") @patch("prkit.core.model_clients.anthropic.os.path.exists") @@ -180,7 +204,7 @@ def test_chat_with_local_image_and_extra_kwargs( ) client = AnthropicModel(ANTHROPIC_TEST_MODEL) - response = client.chat( + response = client.response( "Describe", image_paths=["/tmp/example.png"], temperature=0.2, diff --git a/tests/prkit/core/model_clients/test_base.py b/tests/prkit/core/model_clients/test_base.py index 52e60ea..f4fdd75 100644 --- a/tests/prkit/core/model_clients/test_base.py +++ b/tests/prkit/core/model_clients/test_base.py @@ -8,7 +8,9 @@ import pytest from pydantic import BaseModel -from prkit.core.model_clients.base import BaseModelClient +from prkit.core.domain import PhysicsProblem +from prkit.core.model_clients.base import DEFAULT_INSTRUCTIONS, BaseModelClient +from prkit.core.model_clients.modes import PhysicsOutputMode class TestBaseModelClient: @@ -19,11 +21,10 @@ def test_cannot_instantiate_base_class(self): with pytest.raises(TypeError): BaseModelClient("test-model") - def test_base_class_has_abstract_chat_method(self): - """Test that BaseModelClient defines abstract chat method.""" - # Check that chat is abstract - assert hasattr(BaseModelClient, "chat") - # Try to create a concrete implementation without chat + def test_base_class_has_abstract_response_method(self): + """Test that BaseModelClient defines an abstract response method.""" + assert hasattr(BaseModelClient, "response") + # A concrete implementation without response() stays abstract. with pytest.raises(TypeError): class IncompleteModel(BaseModelClient): @@ -35,7 +36,7 @@ def test_base_class_initialization_attributes(self): """Test that BaseModelClient initializes with correct attributes.""" class ConcreteModel(BaseModelClient): - def chat(self, user_prompt, image_paths=None): + def response(self, input, image_paths=None, **kwargs: Any): return "response" with patch("prkit.core.model_clients.base.load_project_dotenv"): @@ -52,7 +53,7 @@ def test_base_class_with_custom_logger(self): custom_logger = logging.getLogger("custom") class ConcreteModel(BaseModelClient): - def chat(self, user_prompt, image_paths=None): + def response(self, input, image_paths=None, **kwargs: Any): return "response" with patch("prkit.core.model_clients.base.load_project_dotenv"): @@ -63,42 +64,65 @@ def test_base_class_loads_project_dotenv(self): """Test that BaseModelClient loads environment variables.""" class ConcreteModel(BaseModelClient): - def chat(self, user_prompt, image_paths=None): + def response(self, input, image_paths=None, **kwargs: Any): return "response" with patch("prkit.core.model_clients.base.load_project_dotenv") as mock_load: ConcreteModel("test-model") mock_load.assert_called_once() - def test_concrete_implementation_must_implement_chat(self): - """Test that concrete implementations must implement chat method.""" + def test_concrete_implementation_must_implement_response(self): + """Test that concrete implementations can call response().""" class ConcreteModel(BaseModelClient): - def chat(self, user_prompt, image_paths=None): + def response(self, input, image_paths=None, **kwargs: Any): return "response" with patch("prkit.core.model_clients.base.load_project_dotenv"): client = ConcreteModel("test-model") - # Should be able to call chat - result = client.chat("test prompt") + result = client.response("test prompt") assert result == "response" - def test_chat_method_signature(self): - """Test that chat method accepts correct parameters.""" + def test_response_method_signature(self): + """Test that response() accepts the expected parameters.""" class ConcreteModel(BaseModelClient): - def chat(self, user_prompt, image_paths=None): - return f"Response to: {user_prompt}" + def response(self, input, image_paths=None, **kwargs: Any): + return f"Response to: {input}" with patch("prkit.core.model_clients.base.load_project_dotenv"): client = ConcreteModel("test-model") - # Test with text only - result = client.chat("Hello") - assert "Hello" in result + assert "Hello" in client.response("Hello") + assert "Hello" in client.response("Hello", image_paths=["image.jpg"]) - # Test with images - result = client.chat("Hello", image_paths=["image.jpg"]) - assert "Hello" in result + def test_deprecated_chat_alias_warns_and_forwards(self): + """The legacy chat() alias should warn and forward to response().""" + + class ConcreteModel(BaseModelClient): + def response( + self, input, image_paths=None, response_format=None, **kwargs: Any + ): + return f"resp:{input}:{image_paths}" + + with patch("prkit.core.model_clients.base.load_project_dotenv"): + client = ConcreteModel("test-model") + with pytest.warns(DeprecationWarning, match="use .response"): + result = client.chat("Hello", image_paths=["a.jpg"]) + assert result == "resp:Hello:['a.jpg']" + + def test_resolve_instructions_default_and_overrides(self): + """_resolve_instructions falls back to the default but honors explicit values.""" + + class ConcreteModel(BaseModelClient): + def response(self, input, image_paths=None, **kwargs: Any): + return "response" + + with patch("prkit.core.model_clients.base.load_project_dotenv"): + client = ConcreteModel("test-model") + assert client._resolve_instructions(None) == DEFAULT_INSTRUCTIONS + assert client._resolve_instructions("custom") == "custom" + # An explicit empty string is preserved (suppresses the system prompt). + assert client._resolve_instructions("") == "" def test_chat_structured_returns_parsed_model(self): class ResponseModel(BaseModel): @@ -107,8 +131,8 @@ class ResponseModel(BaseModel): class ConcreteModel(BaseModelClient): supports_response_format_json_schema = True - def chat( - self, user_prompt, image_paths=None, response_format=None, **kwargs: Any + def response( + self, input, image_paths=None, response_format=None, **kwargs: Any ): assert response_format["type"] == "json_schema" assert response_format["name"] == "ResponseModel" @@ -127,6 +151,47 @@ def chat( assert result.structured_output_mode == "json_schema" assert result.structured_output_strategy == "dummy_json_schema" + def test_chat_structured_propagates_instructions(self): + """chat_structured should forward `instructions` to response().""" + + class ResponseModel(BaseModel): + answer: str + + seen: dict[str, Any] = {} + + class ConcreteModel(BaseModelClient): + supports_response_format_json_schema = True + + def response( + self, + input, + image_paths=None, + response_format=None, + *, + instructions=None, + **kwargs: Any, + ): + seen["instructions"] = instructions + return '{"answer":"ok"}' + + with patch("prkit.core.model_clients.base.load_project_dotenv"): + client = ConcreteModel("test-model") + client.provider = "dummy" + client.chat_structured( + "Hello", + response_model=ResponseModel, + structured_policy="best_effort", + instructions="be terse", + ) + assert seen["instructions"] == "be terse" + + client.chat_structured( + "Hello", + response_model=ResponseModel, + structured_policy="best_effort", + ) + assert seen["instructions"] is None + def test_chat_structured_native_required_rejects_non_native_provider(self): class ResponseModel(BaseModel): answer: str @@ -135,8 +200,8 @@ class ConcreteModel(BaseModelClient): supports_response_format_json_schema = False supports_response_format_json_object = True - def chat( - self, user_prompt, image_paths=None, response_format=None, **kwargs: Any + def response( + self, input, image_paths=None, response_format=None, **kwargs: Any ): return '{"answer":"ok"}' @@ -151,3 +216,93 @@ def chat( response_model=ResponseModel, structured_policy="native_required", ) + + +class _RecordingModel(BaseModelClient): + """Concrete client that records the arguments passed to response().""" + + def __init__(self, model: str = "test-model"): + with patch("prkit.core.model_clients.base.load_project_dotenv"): + super().__init__(model) + self.calls: list[dict[str, Any]] = [] + + def response( + self, + input, + image_paths=None, + response_format=None, + *, + instructions=None, + **kwargs: Any, + ): + self.calls.append( + { + "input": input, + "image_paths": image_paths, + "response_format": response_format, + "instructions": instructions, + } + ) + return "ANSWER" + + +class TestSolvePhysicsProblem: + """Tests for the high-level solve_physics_problem dispatch.""" + + def test_plain_string_input_passes_through_without_images(self): + client = _RecordingModel() + result = client.solve_physics_problem(" What is g? ") + assert result == "ANSWER" + assert client.calls[0]["input"] == "What is g?" + assert client.calls[0]["image_paths"] is None + + def test_physics_problem_input_builds_prompt_and_attaches_images(self, tmp_path): + image = tmp_path / "fig.png" + image.write_bytes(b"x") + problem = PhysicsProblem( + problem_id="p1", + question="Find the net force.", + problem_type="MC", + options=["1 N", "2 N"], + image_path=[str(image)], + ) + client = _RecordingModel() + client.solve_physics_problem(problem) + + sent = client.calls[0] + assert "p1" in sent["input"] + assert "Find the net force." in sent["input"] + assert "A. 1 N" in sent["input"] and "B. 2 N" in sent["input"] + assert "Attached images: 1 image" in sent["input"] + assert sent["image_paths"] == [str(image)] + + def test_physics_problem_without_images_sends_none(self): + problem = PhysicsProblem(problem_id="p2", question="Define momentum.") + client = _RecordingModel() + client.solve_physics_problem(problem) + assert client.calls[0]["image_paths"] is None + + def test_instructions_forwarded(self): + client = _RecordingModel() + client.solve_physics_problem("Q?", instructions="Answer in one word.") + assert client.calls[0]["instructions"] == "Answer in one word." + + def test_answer_semantics_output_not_implemented(self): + client = _RecordingModel() + with pytest.raises(NotImplementedError): + client.solve_physics_problem( + "Q?", output_mode=PhysicsOutputMode.ANSWER_SEMANTICS + ) + + def test_question_semantics_input_not_implemented(self): + from prkit.semantics.schema import PhysicsQuestionSemantics + + client = _RecordingModel() + qs = PhysicsQuestionSemantics() + with pytest.raises(NotImplementedError, match="Question-semantics"): + client.solve_physics_problem(qs) + + def test_unsupported_input_type_raises_type_error(self): + client = _RecordingModel() + with pytest.raises(TypeError, match="Unsupported problem input type"): + client.solve_physics_problem(123) diff --git a/tests/prkit/core/model_clients/test_dashscope.py b/tests/prkit/core/model_clients/test_dashscope.py index 7044ca3..7503d29 100644 --- a/tests/prkit/core/model_clients/test_dashscope.py +++ b/tests/prkit/core/model_clients/test_dashscope.py @@ -7,6 +7,7 @@ import pytest from pydantic import BaseModel +from prkit.core.model_clients.base import DEFAULT_INSTRUCTIONS from prkit.core.model_clients.dashscope import ( DEFAULT_DASHSCOPE_TIMEOUT_SECONDS, DashscopeModel, @@ -17,6 +18,7 @@ ) DASHSCOPE_TEST_MODEL = "qwen3.6-plus" +SYSTEM_MESSAGE = {"role": "system", "content": DEFAULT_INSTRUCTIONS} class TestDashscopeModel: @@ -89,12 +91,12 @@ def test_chat_text_only(self, _mock_load_project_dotenv, mock_openai_class): mock_client.chat.completions.create.return_value = mock_response client = DashscopeModel(DASHSCOPE_TEST_MODEL) - response = client.chat("Hello, world!") + response = client.response("Hello, world!") assert response == "Test response" mock_client.chat.completions.create.assert_called_once_with( model=DASHSCOPE_TEST_MODEL, - messages=[{"role": "user", "content": "Hello, world!"}], + messages=[SYSTEM_MESSAGE, {"role": "user", "content": "Hello, world!"}], extra_body={"enable_thinking": False}, ) @@ -119,7 +121,7 @@ class ExampleResponse(BaseModel): mock_client.chat.completions.create.return_value = mock_response client = DashscopeModel(DASHSCOPE_TEST_MODEL) - response = client.chat( + response = client.response( "Return JSON only.", response_format=ExampleResponse, max_output_tokens=256, @@ -129,7 +131,8 @@ class ExampleResponse(BaseModel): call_kwargs = mock_client.chat.completions.create.call_args.kwargs assert call_kwargs["model"] == DASHSCOPE_TEST_MODEL assert call_kwargs["messages"] == [ - {"role": "user", "content": "Return JSON only."} + SYSTEM_MESSAGE, + {"role": "user", "content": "Return JSON only."}, ] assert call_kwargs["extra_body"] == {"enable_thinking": False} assert "max_tokens" not in call_kwargs @@ -155,7 +158,7 @@ def test_chat_respects_explicit_extra_body( mock_client.chat.completions.create.return_value = mock_response client = DashscopeModel(DASHSCOPE_TEST_MODEL) - response = client.chat( + response = client.response( "Hello, world!", extra_body={"enable_thinking": True, "thinking_budget": 2048}, ) @@ -163,7 +166,7 @@ def test_chat_respects_explicit_extra_body( assert response == "Test response" mock_client.chat.completions.create.assert_called_once_with( model=DASHSCOPE_TEST_MODEL, - messages=[{"role": "user", "content": "Hello, world!"}], + messages=[SYSTEM_MESSAGE, {"role": "user", "content": "Hello, world!"}], extra_body={"enable_thinking": True, "thinking_budget": 2048}, ) @@ -179,7 +182,7 @@ class ExampleResponse(BaseModel): client = DashscopeModel(DASHSCOPE_TEST_MODEL) with pytest.raises(ValueError, match="not supported in thinking mode"): - client.chat( + client.response( "Return JSON only.", response_format=ExampleResponse, extra_body={"enable_thinking": True}, @@ -194,7 +197,7 @@ def test_chat_rejects_non_dict_extra_body( client = DashscopeModel("custom-model") with pytest.raises(TypeError, match="must be a dict"): - client.chat("Hello", extra_body="bad") + client.response("Hello", extra_body="bad") @patch("prkit.core.model_clients.openai_compatible_chat.OpenAI") @patch("prkit.core.model_clients.base.load_project_dotenv") @@ -215,11 +218,11 @@ def test_default_enable_thinking_respects_env_override( clear=True, ): client = DashscopeModel("custom-model") - response = client.chat("Hello") + response = client.response("Hello") assert response == "ok" mock_client.chat.completions.create.assert_called_once_with( model="custom-model", - messages=[{"role": "user", "content": "Hello"}], + messages=[SYSTEM_MESSAGE, {"role": "user", "content": "Hello"}], extra_body={"enable_thinking": True}, ) diff --git a/tests/prkit/core/model_clients/test_deepseek.py b/tests/prkit/core/model_clients/test_deepseek.py index ad616d5..a19bc04 100644 --- a/tests/prkit/core/model_clients/test_deepseek.py +++ b/tests/prkit/core/model_clients/test_deepseek.py @@ -6,10 +6,12 @@ from pydantic import BaseModel +from prkit.core.model_clients.base import DEFAULT_INSTRUCTIONS from prkit.core.model_clients.deepseek import DeepseekModel DEEPSEEK_CHAT_TEST_MODEL = "deepseek-chat" DEEPSEEK_REASONER_TEST_MODEL = "deepseek-reasoner" +SYSTEM_MESSAGE = {"role": "system", "content": DEFAULT_INSTRUCTIONS} class TestDeepseekModel: @@ -56,12 +58,12 @@ def test_chat_text_only(self, mock_openai_class): mock_client.chat.completions.create.return_value = mock_response client = DeepseekModel(DEEPSEEK_REASONER_TEST_MODEL) - response = client.chat("Hello, world!") + response = client.response("Hello, world!") assert response == "Test response" mock_client.chat.completions.create.assert_called_once_with( model=DEEPSEEK_REASONER_TEST_MODEL, - messages=[{"role": "user", "content": "Hello, world!"}], + messages=[SYSTEM_MESSAGE, {"role": "user", "content": "Hello, world!"}], ) @patch("prkit.core.model_clients.openai_compatible_chat.OpenAI") @@ -79,7 +81,7 @@ def test_chat_with_images_warning(self, mock_openai_class): client = DeepseekModel(DEEPSEEK_CHAT_TEST_MODEL) with patch.object(client.logger, "warning") as mock_warning: - response = client.chat("Hello", image_paths=["image.jpg"]) + response = client.response("Hello", image_paths=["image.jpg"]) assert response == "Response" mock_warning.assert_called_once() @@ -99,11 +101,12 @@ def test_chat_ignores_images(self, mock_openai_class): mock_client.chat.completions.create.return_value = mock_response client = DeepseekModel(DEEPSEEK_REASONER_TEST_MODEL) - client.chat("Hello", image_paths=["image1.jpg", "image2.png"]) + client.response("Hello", image_paths=["image1.jpg", "image2.png"]) call_kwargs = mock_client.chat.completions.create.call_args[1] - assert len(call_kwargs["messages"]) == 1 - assert call_kwargs["messages"][0]["content"] == "Hello" + assert len(call_kwargs["messages"]) == 2 + assert call_kwargs["messages"][0] == SYSTEM_MESSAGE + assert call_kwargs["messages"][1]["content"] == "Hello" @patch("prkit.core.model_clients.openai_compatible_chat.OpenAI") def test_chat_response_format_uses_json_object(self, mock_openai_class): @@ -123,7 +126,7 @@ class ExampleResponse(BaseModel): mock_client.chat.completions.create.return_value = mock_response client = DeepseekModel(DEEPSEEK_CHAT_TEST_MODEL) - response = client.chat( + response = client.response( "Return JSON only.", response_format=ExampleResponse, max_output_tokens=512, @@ -134,5 +137,6 @@ class ExampleResponse(BaseModel): assert call_kwargs["model"] == DEEPSEEK_CHAT_TEST_MODEL assert call_kwargs["response_format"] == {"type": "json_object"} assert call_kwargs["max_tokens"] == 512 - assert "Return JSON only." in call_kwargs["messages"][0]["content"] - assert "Return ONLY JSON" in call_kwargs["messages"][0]["content"] + assert call_kwargs["messages"][0] == SYSTEM_MESSAGE + assert "Return JSON only." in call_kwargs["messages"][1]["content"] + assert "Return ONLY JSON" in call_kwargs["messages"][1]["content"] diff --git a/tests/prkit/core/model_clients/test_gemini.py b/tests/prkit/core/model_clients/test_gemini.py index c7263d7..0cf4dfb 100644 --- a/tests/prkit/core/model_clients/test_gemini.py +++ b/tests/prkit/core/model_clients/test_gemini.py @@ -8,6 +8,7 @@ import pytest from pydantic import BaseModel +from prkit.core.model_clients.base import DEFAULT_INSTRUCTIONS from prkit.core.model_clients.gemini import ( GeminiModel, _extract_gemini_error_details, @@ -69,7 +70,7 @@ def test_chat_text_only(self, mock_genai): mock_client.models.generate_content.return_value = mock_response client = GeminiModel(GEMINI_TEST_MODEL) - response = client.chat("Hello, world!") + response = client.response("Hello, world!") assert response == "Test response" mock_client.models.generate_content.assert_called_once() @@ -81,6 +82,27 @@ def test_chat_text_only(self, mock_genai): config = call_kwargs["config"] assert config is not None assert config.max_output_tokens == 65535 + # Default system prompt is sent via Gemini's system_instruction config. + assert config.system_instruction == DEFAULT_INSTRUCTIONS + + @patch("prkit.core.model_clients.gemini.genai") + def test_response_system_instruction_explicit_and_suppressed(self, mock_genai): + """Explicit instructions override the default; empty string suppresses it.""" + mock_client = MagicMock() + mock_genai.Client.return_value = mock_client + mock_response = Mock() + mock_response.text = "ok" + mock_client.models.generate_content.return_value = mock_response + + client = GeminiModel(GEMINI_TEST_MODEL) + + client.response("Hi", instructions="Custom system.") + config = mock_client.models.generate_content.call_args[1]["config"] + assert config.system_instruction == "Custom system." + + client.response("Hi", instructions="") + config = mock_client.models.generate_content.call_args[1]["config"] + assert config.system_instruction is None @patch("prkit.core.model_clients.gemini.genai") def test_chat_with_kwargs(self, mock_genai): @@ -92,7 +114,7 @@ def test_chat_with_kwargs(self, mock_genai): mock_client.models.generate_content.return_value = mock_response client = GeminiModel(GEMINI_TEST_MODEL) - client.chat("Hello", temperature=0.7) + client.response("Hello", temperature=0.7) call_kwargs = mock_client.models.generate_content.call_args[1] assert "config" in call_kwargs @@ -110,7 +132,7 @@ def test_chat_without_config_kwargs(self, mock_genai): mock_client.models.generate_content.return_value = mock_response client = GeminiModel(GEMINI_TEST_MODEL) - client.chat("Hello") + client.response("Hello") call_kwargs = mock_client.models.generate_content.call_args[1] config = call_kwargs["config"] @@ -128,7 +150,7 @@ def test_chat_with_images_error(self, mock_genai): client = GeminiModel(GEMINI_TEST_MODEL) with patch.object(client.logger, "error") as mock_error: - response = client.chat("Hello", image_paths=["image.jpg"]) + response = client.response("Hello", image_paths=["image.jpg"]) assert response == "Response" mock_error.assert_called_once() @@ -165,7 +187,7 @@ class ExampleResponse(BaseModel): mock_client.models.generate_content.return_value = mock_response client = GeminiModel(GEMINI_TEST_MODEL) - response = client.chat( + response = client.response( "Return JSON", image_paths=["/tmp/example.png"], response_format=ExampleResponse, @@ -196,7 +218,7 @@ def test_chat_logs_failed_image_open( client = GeminiModel(GEMINI_TEST_MODEL) with patch.object(client.logger, "error") as mock_error: - response = client.chat("Hello", image_paths=["/tmp/bad.png"]) + response = client.response("Hello", image_paths=["/tmp/bad.png"]) assert response == "Response" mock_error.assert_called_once() @@ -218,7 +240,7 @@ def test_chat_empty_response_raises_runtime_error_with_details(self, mock_genai) RuntimeError, match="prompt_block_reason=SAFETY; finish_reason=RECITATION", ): - client.chat("Hello") + client.response("Hello") def test_extract_gemini_error_details_handles_empty_and_prompt_blocks(self): empty = SimpleNamespace(prompt_feedback=None, candidates=None) diff --git a/tests/prkit/core/model_clients/test_ollama.py b/tests/prkit/core/model_clients/test_ollama.py index 1df3534..5b7cbd4 100644 --- a/tests/prkit/core/model_clients/test_ollama.py +++ b/tests/prkit/core/model_clients/test_ollama.py @@ -8,10 +8,12 @@ import pytest from pydantic import BaseModel +from prkit.core.model_clients.base import DEFAULT_INSTRUCTIONS from prkit.core.model_clients.ollama import OllamaModel, normalize_ollama_model_name OLLAMA_QWEN_TEST_MODEL = "ollama/qwen3.5:397b-cloud" OLLAMA_MISTRAL_TEST_MODEL = "ollama/mistral-large-3:675b-cloud" +SYSTEM_MESSAGE = {"role": "system", "content": DEFAULT_INSTRUCTIONS} class TestOllamaModel: @@ -117,15 +119,16 @@ def test_chat_text_only(self, mock_ollama_module): mock_ollama_module.chat.return_value = mock_response client = OllamaModel(OLLAMA_QWEN_TEST_MODEL) - response = client.chat("Hello, world!") + response = client.response("Hello, world!") assert response == "Test response" mock_ollama_module.chat.assert_called_once() call_kwargs = mock_ollama_module.chat.call_args[1] assert call_kwargs["model"] == "qwen3.5:397b-cloud" - assert len(call_kwargs["messages"]) == 1 - assert call_kwargs["messages"][0]["role"] == "user" - assert call_kwargs["messages"][0]["content"] == "Hello, world!" + assert len(call_kwargs["messages"]) == 2 + assert call_kwargs["messages"][0] == SYSTEM_MESSAGE + assert call_kwargs["messages"][1]["role"] == "user" + assert call_kwargs["messages"][1]["content"] == "Hello, world!" assert "format" not in call_kwargs @patch("prkit.core.model_clients.ollama.ollama") @@ -148,7 +151,7 @@ class ExampleResponse(BaseModel): mock_ollama_module.chat.return_value = mock_response client = OllamaModel(OLLAMA_QWEN_TEST_MODEL) - response = client.chat( + response = client.response( "Return JSON only.", response_format=ExampleResponse, max_output_tokens=256, @@ -177,7 +180,7 @@ def test_chat_with_json_object_response_format_uses_json_mode( mock_ollama_module.chat.return_value = mock_response client = OllamaModel(OLLAMA_QWEN_TEST_MODEL) - response = client.chat( + response = client.response( "Return JSON only.", response_format={"type": "json_object"}, ) @@ -199,7 +202,7 @@ def test_chat_uses_normalized_prefixed_model_name(self, mock_ollama_module): mock_ollama_module.chat.return_value = mock_response client = OllamaModel(OLLAMA_QWEN_TEST_MODEL) - response = client.chat("Hello") + response = client.response("Hello") assert response == "Response" call_kwargs = mock_ollama_module.chat.call_args[1] @@ -224,14 +227,15 @@ def test_chat_with_images(self, mock_ollama_module, tmp_path): mock_ollama_module.chat.return_value = mock_response client = OllamaModel(OLLAMA_QWEN_TEST_MODEL) - response = client.chat( + response = client.response( "Describe these images", image_paths=[str(image1), str(image2)] ) assert response == "Image description" call_kwargs = mock_ollama_module.chat.call_args[1] - assert "images" in call_kwargs["messages"][0] - assert len(call_kwargs["messages"][0]["images"]) == 2 + assert call_kwargs["messages"][0] == SYSTEM_MESSAGE + assert "images" in call_kwargs["messages"][1] + assert len(call_kwargs["messages"][1]["images"]) == 2 @patch("prkit.core.model_clients.ollama.ollama") def test_chat_with_nonexistent_image(self, mock_ollama_module): @@ -242,7 +246,7 @@ def test_chat_with_nonexistent_image(self, mock_ollama_module): client = OllamaModel(OLLAMA_MISTRAL_TEST_MODEL) with pytest.raises(FileNotFoundError, match="Image file not found"): - client.chat("Describe this", image_paths=["/nonexistent/image.jpg"]) + client.response("Describe this", image_paths=["/nonexistent/image.jpg"]) @patch("prkit.core.model_clients.ollama.ollama") def test_chat_with_base_url(self, mock_ollama_module): @@ -253,7 +257,7 @@ def test_chat_with_base_url(self, mock_ollama_module): mock_ollama_module.Client.return_value = mock_client client = OllamaModel(OLLAMA_MISTRAL_TEST_MODEL, base_url="http://custom:11434") - response = client.chat("Hello") + response = client.response("Hello") assert response == "Response" mock_client.chat.assert_called_once() @@ -271,7 +275,7 @@ def test_chat_model_not_found_error(self, mock_ollama_module): client = OllamaModel("unknown-model") with pytest.raises(ValueError, match="Model 'unknown-model' not found"): - client.chat("Hello") + client.response("Hello") @patch("prkit.core.model_clients.ollama.ollama") def test_chat_connection_error(self, mock_ollama_module): @@ -285,7 +289,7 @@ def test_chat_connection_error(self, mock_ollama_module): client = OllamaModel(OLLAMA_QWEN_TEST_MODEL) with pytest.raises(ConnectionError, match="Ollama service is not running"): - client.chat("Hello") + client.response("Hello") @patch("prkit.core.model_clients.ollama.ollama") def test_chat_response_dict_format(self, mock_ollama_module): @@ -298,7 +302,7 @@ def test_chat_response_dict_format(self, mock_ollama_module): mock_ollama_module.chat.return_value = mock_response client = OllamaModel(OLLAMA_MISTRAL_TEST_MODEL) - response = client.chat("Hello") + response = client.response("Hello") assert response == "Dict response" @@ -315,7 +319,7 @@ def test_chat_with_empty_image_list(self, mock_ollama_module): mock_ollama_module.chat.return_value = mock_response client = OllamaModel(OLLAMA_QWEN_TEST_MODEL) - response = client.chat("Hello", image_paths=[]) + response = client.response("Hello", image_paths=[]) assert response == "Response" call_kwargs = mock_ollama_module.chat.call_args[1] @@ -334,7 +338,7 @@ def test_chat_temperature_option(self, mock_ollama_module): mock_ollama_module.chat.return_value = mock_response client = OllamaModel(OLLAMA_QWEN_TEST_MODEL) - client.chat("Hello") + client.response("Hello") call_kwargs = mock_ollama_module.chat.call_args[1] assert "options" in call_kwargs @@ -405,7 +409,7 @@ def test_no_api_key_falls_back_to_module_level_chat(self, mock_ollama_module): mock_ollama_module.chat.return_value = mock_response client = OllamaModel(OLLAMA_QWEN_TEST_MODEL) - client.chat("Hello") + client.response("Hello") mock_ollama_module.chat.assert_called_once() diff --git a/tests/prkit/core/model_clients/test_openai.py b/tests/prkit/core/model_clients/test_openai.py index 2bf033e..f07b1a0 100644 --- a/tests/prkit/core/model_clients/test_openai.py +++ b/tests/prkit/core/model_clients/test_openai.py @@ -100,7 +100,7 @@ def test_chat_text_only(self, mock_openai_class): mock_client.responses.create.return_value = mock_response client = OpenAIModel(OPENAI_TEST_MODEL) - response = client.chat("Hello, world!") + response = client.response("Hello, world!") assert response == "Test response" mock_client.responses.create.assert_called_once() @@ -111,6 +111,28 @@ def test_chat_text_only(self, mock_openai_class): assert len(call_kwargs["input"][0]["content"]) == 1 assert call_kwargs["input"][0]["content"][0]["type"] == "input_text" assert call_kwargs["input"][0]["content"][0]["text"] == "Hello, world!" + # OpenAI works with input alone — no default system prompt. + assert "instructions" not in call_kwargs + + @patch("prkit.core.model_clients.openai.OpenAI") + def test_response_instructions_only_when_provided(self, mock_openai_class): + """OpenAI sends `instructions` only when the caller supplies it.""" + mock_client = MagicMock() + mock_openai_class.return_value = mock_client + mock_response = Mock() + mock_response.output_text = "ok" + mock_client.responses.create.return_value = mock_response + + client = OpenAIModel(OPENAI_TEST_MODEL) + + client.response("Hi") + assert "instructions" not in mock_client.responses.create.call_args[1] + + client.response("Hi", instructions="Answer tersely.") + assert ( + mock_client.responses.create.call_args[1]["instructions"] + == "Answer tersely." + ) @patch("prkit.core.model_clients.openai.OpenAI") def test_chat_with_images(self, mock_openai_class, tmp_path): @@ -127,7 +149,7 @@ def test_chat_with_images(self, mock_openai_class, tmp_path): mock_client.responses.create.return_value = mock_response client = OpenAIModel(OPENAI_TEST_MODEL) - response = client.chat("Describe this image", image_paths=[str(image_file)]) + response = client.response("Describe this image", image_paths=[str(image_file)]) assert response == "Image description" call_kwargs = mock_client.responses.create.call_args[1] @@ -151,7 +173,7 @@ def test_chat_with_http_url(self, mock_openai_class): mock_client.responses.create.return_value = mock_response client = OpenAIModel(OPENAI_TEST_MODEL) - response = client.chat( + response = client.response( "Describe this image", image_paths=["https://example.com/image.jpg"] ) @@ -174,7 +196,7 @@ def test_chat_with_base64_data_url(self, mock_openai_class): client = OpenAIModel(OPENAI_TEST_MODEL) data_url = "data:image/jpeg;base64,/9j/4AAQSkZJRg==" - response = client.chat("Describe this image", image_paths=[data_url]) + response = client.response("Describe this image", image_paths=[data_url]) assert response == "Base64 image description" call_kwargs = mock_client.responses.create.call_args[1] @@ -191,7 +213,7 @@ def test_chat_o_family_with_reasoning(self, mock_openai_class): mock_client.responses.create.return_value = mock_response client = OpenAIModel("o3") - client.chat("Solve this problem") + client.response("Solve this problem") call_kwargs = mock_client.responses.create.call_args[1] assert "reasoning" in call_kwargs @@ -207,7 +229,7 @@ def test_chat_non_o_family_no_reasoning(self, mock_openai_class): mock_client.responses.create.return_value = mock_response client = OpenAIModel(OPENAI_TEST_MODEL) - client.chat("Hello") + client.response("Hello") call_kwargs = mock_client.responses.create.call_args[1] assert "reasoning" not in call_kwargs @@ -222,7 +244,7 @@ def test_chat_forwards_max_output_tokens(self, mock_openai_class): mock_client.responses.create.return_value = mock_response client = OpenAIModel(OPENAI_TEST_MODEL) - client.chat("Hello", max_output_tokens=321) + client.response("Hello", max_output_tokens=321) call_kwargs = mock_client.responses.create.call_args[1] assert call_kwargs["max_output_tokens"] == 321 @@ -242,7 +264,7 @@ class ExampleResponse(BaseModel): mock_client.responses.create.return_value = mock_response client = OpenAIModel(OPENAI_TEST_MODEL) - client.chat("Hello", response_format=ExampleResponse) + client.response("Hello", response_format=ExampleResponse) schema = mock_client.responses.create.call_args[1]["text"]["format"]["schema"] assert set(schema["required"]) == set(schema["properties"].keys()) @@ -370,7 +392,7 @@ def test_chat_with_multiple_images(self, mock_openai_class, tmp_path): mock_client.responses.create.return_value = mock_response client = OpenAIModel(OPENAI_TEST_MODEL) - response = client.chat("Describe these images", image_paths=images) + response = client.response("Describe these images", image_paths=images) assert response == "Multi-image response" call_kwargs = mock_client.responses.create.call_args[1] @@ -388,7 +410,7 @@ def test_chat_with_empty_string_prompt(self, mock_openai_class): mock_client.responses.create.return_value = mock_response client = OpenAIModel(OPENAI_TEST_MODEL) - response = client.chat("") + response = client.response("") assert response == "Response" call_kwargs = mock_client.responses.create.call_args[1] @@ -404,7 +426,7 @@ def test_chat_with_none_image_paths(self, mock_openai_class): mock_client.responses.create.return_value = mock_response client = OpenAIModel(OPENAI_TEST_MODEL) - response = client.chat("Hello", image_paths=None) + response = client.response("Hello", image_paths=None) assert response == "Response" call_kwargs = mock_client.responses.create.call_args[1] diff --git a/tests/prkit/core/model_clients/test_openai_compatible_chat.py b/tests/prkit/core/model_clients/test_openai_compatible_chat.py index 6aefeaf..113dd11 100644 --- a/tests/prkit/core/model_clients/test_openai_compatible_chat.py +++ b/tests/prkit/core/model_clients/test_openai_compatible_chat.py @@ -2,10 +2,13 @@ from pydantic import BaseModel +from prkit.core.model_clients.base import DEFAULT_INSTRUCTIONS from prkit.core.model_clients.openai_compatible_chat import ( OpenAICompatibleChatModel, ) +SYSTEM_MESSAGE = {"role": "system", "content": DEFAULT_INSTRUCTIONS} + class DummyChatProvider(OpenAICompatibleChatModel): provider_name = "dummy" @@ -91,7 +94,7 @@ class ExampleResponse(BaseModel): with patch.dict("os.environ", {"DUMMY_API_KEY": "test-key"}, clear=True): client = DummyChatProvider("model-a") - response = client.chat( + response = client.response( "Return JSON only.", response_format=ExampleResponse, max_output_tokens=123, @@ -102,7 +105,8 @@ class ExampleResponse(BaseModel): call_kwargs = mock_client.chat.completions.create.call_args[1] assert call_kwargs["model"] == "model-a" assert call_kwargs["messages"] == [ - {"role": "user", "content": "Return JSON only."} + SYSTEM_MESSAGE, + {"role": "user", "content": "Return JSON only."}, ] assert call_kwargs["max_tokens"] == 123 assert call_kwargs["temperature"] == 0.2 diff --git a/tests/prkit/core/model_clients/test_package_init.py b/tests/prkit/core/model_clients/test_package_init.py index 17a285f..2a51067 100644 --- a/tests/prkit/core/model_clients/test_package_init.py +++ b/tests/prkit/core/model_clients/test_package_init.py @@ -32,3 +32,18 @@ def test_model_clients_getattr_returns_lazy_classes(): def test_model_clients_getattr_rejects_unknown_names(): with pytest.raises(AttributeError): model_clients.__getattr__("UnknownModel") + + +def test_new_public_symbols_are_exported(): + """DEFAULT_INSTRUCTIONS, PhysicsOutputMode, and prompt helpers are public.""" + for name in ( + "DEFAULT_INSTRUCTIONS", + "PhysicsOutputMode", + "build_plain_question_prompt", + "format_problem_context", + ): + assert name in model_clients.__all__ + assert getattr(model_clients, name) is not None + + assert isinstance(model_clients.DEFAULT_INSTRUCTIONS, str) + assert model_clients.DEFAULT_INSTRUCTIONS.strip() diff --git a/tests/prkit/core/model_clients/test_prompts.py b/tests/prkit/core/model_clients/test_prompts.py new file mode 100644 index 0000000..6a7581e --- /dev/null +++ b/tests/prkit/core/model_clients/test_prompts.py @@ -0,0 +1,64 @@ +"""Tests for the core-layer physics prompt builders.""" + +from prkit.core.domain import PhysicsProblem +from prkit.core.model_clients.prompts import ( + build_plain_question_prompt, + format_problem_context, +) + + +def test_format_problem_context_includes_core_fields(): + problem = PhysicsProblem( + problem_id="p1", + question=" What is the net force? ", + problem_type="MC", + domain="mechanics", + options=["1 N", "2 N", "3 N"], + ) + text = format_problem_context(problem) + + assert "Problem ID: p1" in text + assert "Problem type: MC" in text + assert "Language: en" in text + assert "A. 1 N" in text and "B. 2 N" in text and "C. 3 N" in text + # Question is stripped and rendered last. + assert text.endswith("Question:\nWhat is the net force?") + + +def test_format_problem_context_notes_attached_images(tmp_path): + image = tmp_path / "fig.png" + image.write_bytes(b"x") + problem = PhysicsProblem( + problem_id="p2", + question="Describe the figure.", + image_path=[str(image)], + ) + text = format_problem_context(problem) + assert "Attached images: 1 image(s)" in text + + +def test_format_problem_context_omits_optional_sections_when_absent(): + problem = PhysicsProblem(problem_id="p3", question="Define momentum.") + text = format_problem_context(problem) + assert "Options:" not in text + assert "Attached images" not in text + + +def test_build_plain_question_prompt_matches_context(): + problem = PhysicsProblem(problem_id="p4", question="State Hooke's law.") + assert build_plain_question_prompt(problem) == format_problem_context(problem) + + +def test_header_parity_with_semantics_format_problem(): + """The core header must match the semantics layer's prediction header.""" + from prkit.semantics.inference.prompts import _format_problem + + problem = PhysicsProblem( + problem_id="p5", + question="A 2 kg block accelerates at 3 m/s^2. Net force?", + problem_type="OE", + options=None, + ) + assert format_problem_context(problem) == _format_problem( + problem, include_reference_context=False + ) diff --git a/tests/prkit/core/model_clients/test_xai.py b/tests/prkit/core/model_clients/test_xai.py index 19c4c92..be374a5 100644 --- a/tests/prkit/core/model_clients/test_xai.py +++ b/tests/prkit/core/model_clients/test_xai.py @@ -6,10 +6,12 @@ from pydantic import BaseModel, Field +from prkit.core.model_clients.base import DEFAULT_INSTRUCTIONS from prkit.core.model_clients.structured_output import coerce_structured_output_spec from prkit.core.model_clients.xai import XAIModel XAI_TEST_MODEL = "grok-4-1-fast-reasoning" +SYSTEM_MESSAGE = {"role": "system", "content": DEFAULT_INSTRUCTIONS} class TestXAIModel: @@ -47,12 +49,12 @@ def test_chat_text_only(self, _mock_load_project_dotenv, mock_openai_class): mock_client.chat.completions.create.return_value = mock_response client = XAIModel(XAI_TEST_MODEL) - response = client.chat("Hello, world!") + response = client.response("Hello, world!") assert response == "Test response" mock_client.chat.completions.create.assert_called_once_with( model=XAI_TEST_MODEL, - messages=[{"role": "user", "content": "Hello, world!"}], + messages=[SYSTEM_MESSAGE, {"role": "user", "content": "Hello, world!"}], ) @patch("prkit.core.model_clients.openai_compatible_chat.OpenAI") diff --git a/tests/prkit/semantics/test_inference_prompts.py b/tests/prkit/semantics/test_inference_prompts.py index eb22ad3..a0afbc0 100644 --- a/tests/prkit/semantics/test_inference_prompts.py +++ b/tests/prkit/semantics/test_inference_prompts.py @@ -39,15 +39,15 @@ def __init__(self) -> None: self.last_prompt: str | None = None self.last_response_format: dict[str, Any] | type | None = None - def chat( + def response( self, - user_prompt: str, + input: str, image_paths: list[str] | None = None, response_format: dict[str, Any] | type | None = None, **kwargs: Any, ) -> str: del image_paths, kwargs - self.last_prompt = user_prompt + self.last_prompt = input self.last_response_format = response_format return json.dumps( { @@ -162,15 +162,15 @@ def __init__(self) -> None: self.last_prompt: str | None = None self.last_response_format: dict[str, Any] | type | None = None - def chat( + def response( self, - user_prompt: str, + input: str, image_paths: list[str] | None = None, response_format: dict[str, Any] | type | None = None, **kwargs: Any, ) -> str: del image_paths, kwargs - self.last_prompt = user_prompt + self.last_prompt = input self.last_response_format = response_format return json.dumps( { @@ -267,15 +267,15 @@ def __init__(self) -> None: self.response_formats: list[dict[str, Any] | type | None] = [] self.max_output_tokens: list[int | None] = [] - def chat( + def response( self, - user_prompt: str, + input: str, image_paths: list[str] | None = None, response_format: dict[str, Any] | type | None = None, **kwargs: Any, ) -> str: del image_paths - self.prompts.append(user_prompt) + self.prompts.append(input) self.response_formats.append(response_format) self.max_output_tokens.append(kwargs.get("max_output_tokens")) if len(self.prompts) == 1: @@ -726,14 +726,14 @@ def __init__(self) -> None: super().__init__(model="stub-model") self.provider = "stub" - def chat( + def response( self, - user_prompt: str, + input: str, image_paths: list[str] | None = None, response_format: dict[str, Any] | type | None = None, **kwargs: Any, ) -> str: - del user_prompt, image_paths, response_format, kwargs + del input, image_paths, response_format, kwargs raise AssertionError( "chat should not be called when native json_schema is unsupported" ) From dff9d014034b753841399a6924c8da89f16c6490 Mon Sep 17 00:00:00 2001 From: Yinghuan Zhang Date: Sun, 14 Jun 2026 13:42:17 -0400 Subject: [PATCH 2/8] Update cookbooks and docs for response() / solve_physics_problem() Switch the model-client examples in the cookbooks, docs/CORE.md, README, and RELEASE_NOTES to the new response(input=...) call, document the `instructions` parameter and DEFAULT_INSTRUCTIONS fallback, and add a solve_physics_problem() usage section. Also fixes a pre-existing typo in the README quickstart. Co-Authored-By: Claude Opus 4.8 (1M context) --- README.md | 2 +- RELEASE_NOTES.md | 2 +- cookbooks/inference_deepseek.py | 4 +- cookbooks/inference_gemini.py | 4 +- cookbooks/inference_ollama.py | 2 +- cookbooks/inference_openai.py | 2 +- cookbooks/inference_seephys_structured.py | 4 +- .../inference_single_with_answer_tags.py | 2 +- docs/CORE.md | 42 +++++++++++++++++-- 9 files changed, 49 insertions(+), 15 deletions(-) diff --git a/README.md b/README.md index a05e941..12cc861 100644 --- a/README.md +++ b/README.md @@ -25,7 +25,7 @@ dataset = DatasetHub.load("physreason", variant="full", split="test") # Run inference with the unified model client (core component) client = create_model_client("gpt-4.1-mini") for problem in dataset[:3]: - print(client.chat(problem.question)[:200) + print(client.solve_physics_problem(problem)[:200]) ``` The same pattern works across different datasets and model providers—swap the dataset name or model identifier. diff --git a/RELEASE_NOTES.md b/RELEASE_NOTES.md index 968c4ee..9a8611e 100644 --- a/RELEASE_NOTES.md +++ b/RELEASE_NOTES.md @@ -66,7 +66,7 @@ from prkit.core.model_clients import create_model_client dataset = DatasetHub.load("physreason", variant="full", split="test") client = create_model_client("gpt-4.1-mini") for problem in dataset[:3]: - print(client.chat(problem.question)[:200]) + print(client.solve_physics_problem(problem)[:200]) ``` ### Requirements diff --git a/cookbooks/inference_deepseek.py b/cookbooks/inference_deepseek.py index 68f2e16..2a2e026 100755 --- a/cookbooks/inference_deepseek.py +++ b/cookbooks/inference_deepseek.py @@ -135,8 +135,8 @@ def run_inference(model_name: str, prompt: str, image_path: str = None): try: # Note: image_paths parameter is accepted but images will be ignored - response = client.chat( - user_prompt=prompt, image_paths=[image_path] if image_path else None + response = client.response( + input=prompt, image_paths=[image_path] if image_path else None ) logger.info("\n" + "=" * 60) diff --git a/cookbooks/inference_gemini.py b/cookbooks/inference_gemini.py index 564d3a4..555ea43 100755 --- a/cookbooks/inference_gemini.py +++ b/cookbooks/inference_gemini.py @@ -151,8 +151,8 @@ def run_inference(model_name: str, prompt: str, image_path: str = None): try: # Note: image_paths parameter is accepted but images will be ignored - response = client.chat( - user_prompt=prompt, image_paths=[image_path] if image_path else None + response = client.response( + input=prompt, image_paths=[image_path] if image_path else None ) logger.info("\n" + "=" * 60) diff --git a/cookbooks/inference_ollama.py b/cookbooks/inference_ollama.py index c7cfc36..6ce33d6 100755 --- a/cookbooks/inference_ollama.py +++ b/cookbooks/inference_ollama.py @@ -134,7 +134,7 @@ def run_inference(model_name: str, prompt: str, image_path: str = None): logger.info(f" Images: {len(image_paths)} image(s)") try: - response = client.chat(user_prompt=prompt, image_paths=image_paths) + response = client.response(input=prompt, image_paths=image_paths) logger.info("\n" + "=" * 60) logger.info("Response:") diff --git a/cookbooks/inference_openai.py b/cookbooks/inference_openai.py index bc3272b..e4da232 100755 --- a/cookbooks/inference_openai.py +++ b/cookbooks/inference_openai.py @@ -141,7 +141,7 @@ def run_inference(model_name: str, prompt: str, image_path: str = None): logger.info(f" Images: {len(image_paths)} image(s)") try: - response = client.chat(user_prompt=prompt, image_paths=image_paths) + response = client.response(input=prompt, image_paths=image_paths) logger.info("\n" + "=" * 60) logger.info("Response:") diff --git a/cookbooks/inference_seephys_structured.py b/cookbooks/inference_seephys_structured.py index dc95675..8c4bd40 100644 --- a/cookbooks/inference_seephys_structured.py +++ b/cookbooks/inference_seephys_structured.py @@ -172,8 +172,8 @@ def run_inference( if client.provider == "google": extra_kwargs["max_output_tokens"] = 65535 - response_text = client.chat( - user_prompt=full_prompt, + response_text = client.response( + input=full_prompt, image_paths=image_paths, response_format=ReasonAndAnswer, **extra_kwargs, diff --git a/cookbooks/inference_single_with_answer_tags.py b/cookbooks/inference_single_with_answer_tags.py index a3eaf4e..b4daa0a 100644 --- a/cookbooks/inference_single_with_answer_tags.py +++ b/cookbooks/inference_single_with_answer_tags.py @@ -175,7 +175,7 @@ def run_single_inference( for p in image_paths: if not Path(p).exists(): raise FileNotFoundError(f"Image file not found: {p}") - response = client.chat(user_prompt=full_prompt, image_paths=image_paths) + response = client.response(input=full_prompt, image_paths=image_paths) result["model_response"] = response result["model_answer"] = parse_answer_from_response(response) except FileNotFoundError as e: diff --git a/docs/CORE.md b/docs/CORE.md index 147afa5..b17152c 100644 --- a/docs/CORE.md +++ b/docs/CORE.md @@ -166,7 +166,11 @@ Utility components provide supporting infrastructure used across the toolkit. Th ### Model Client (BaseModelClient, create_model_client) -Unified interface for running inference across multiple providers (LLMs and VLMs). Subclasses implement `chat(user_prompt: str, image_paths: Optional[List[str]] = None)`. Use `create_model_client(model: str)` to get the right implementation based on the model name. Vision-capable providers consume `image_paths`; others ignore images with a warning. +Unified interface for running inference across multiple providers (LLMs and VLMs). Subclasses implement `response(input: str, image_paths: Optional[List[str]] = None, *, instructions: Optional[str] = None)`, mirroring OpenAI's `client.responses.create` (`input` is the user prompt, `instructions` is the system prompt). Use `create_model_client(model: str)` to get the right implementation based on the model name. Vision-capable providers consume `image_paths`; others ignore images with a warning. + +When `instructions` is omitted, every provider **except OpenAI** falls back to a short default system prompt, `DEFAULT_INSTRUCTIONS` (`"You are a physics expert. …"`); OpenAI sends `input` alone. Pass `instructions=""` to suppress the system prompt entirely. The legacy `chat(user_prompt=...)` method still works as a deprecated alias for `response(input=...)` but emits a `DeprecationWarning`. + +For physics problems specifically, `solve_physics_problem()` builds the prompt and attaches images for you (see below). **Supported providers** (selected by model name pattern): @@ -185,14 +189,44 @@ Unified interface for running inference across multiple providers (LLMs and VLMs from prkit.core.model_clients import create_model_client client = create_model_client("gpt-4.1-mini") -print(client.chat("State Newton's second law in one sentence.")) +print(client.response("State Newton's second law in one sentence.")) # Vision (optional) -text = client.chat( +text = client.response( "Solve the problem shown in the image and return only the final answer.", image_paths=["/absolute/path/to/problem.png"], ) print(text) + +# Custom system prompt (sent as the provider's system/instructions field) +print(client.response("List three SI base units.", instructions="Answer tersely.")) +``` + +#### Asking a physics problem (`solve_physics_problem`) + +`solve_physics_problem()` is a convenience that builds the prompt and attaches any +images, then calls `response()`. The input is dispatched on type: a plain `str` +question, a `PhysicsProblem` (parsed into prompt text plus its `image_path` +images), or — in a future release — a `PhysicsQuestionSemantics`. The +`output_mode` selects the answer form; only `PhysicsOutputMode.ANSWER_TEXT` is +implemented today. + +```python +from prkit.core.domain import PhysicsProblem +from prkit.core.model_clients import create_model_client + +client = create_model_client("gpt-4.1-mini") + +# From a plain question string +print(client.solve_physics_problem("State Newton's second law in one sentence.")) + +# From a PhysicsProblem (question + options + images are formatted for you) +problem = PhysicsProblem( + problem_id="p1", + question="A 2 kg block accelerates at 3 m/s^2. What net force acts on it?", + problem_type="OE", +) +print(client.solve_physics_problem(problem)) ``` #### Custom OpenAI Responses-API endpoints @@ -248,7 +282,7 @@ client = OllamaModel("llama3:70b-cloud", base_url="https://ollama.com") Key-resolution precedence: explicit `api_key` → `api_key_env` env lookup → library auto-reads `OLLAMA_API_KEY`. For remote hosts (`base_url` pointing to a non-localhost address) a failed startup preflight emits a warning instead of raising `ConnectionError`; -precise errors surface at `chat()` call time. +precise errors surface at `response()` call time. #### Registering additional providers From 2563141c9522cee22ca51815862911266ca636dd Mon Sep 17 00:00:00 2001 From: Yinghuan Zhang Date: Mon, 15 Jun 2026 08:52:36 -0400 Subject: [PATCH 3/8] Add Batch API support across OpenAI, Anthropic, and Gemini Add a synchronous batch job lifecycle (submit_batch/poll_batch/ retrieve_batch_results) and a free-text build_batch_request() to BaseModelClient, complementing the existing structured batch builder. New batch_types module (BatchState/BatchStatus/BatchItemStatus/ BatchResult) normalizes each provider's status enum and per-request results. Each provider's request-body construction is now shared between response() and the batch builders to prevent drift; o-family OpenAI drops temperature at build time. Co-Authored-By: Claude Opus 4.8 (1M context) --- CHANGELOG.md | 1 + src/prkit/core/model_clients/__init__.py | 12 + src/prkit/core/model_clients/anthropic.py | 300 ++++++++++------ src/prkit/core/model_clients/base.py | 87 ++++- src/prkit/core/model_clients/batch_types.py | 89 +++++ src/prkit/core/model_clients/gemini.py | 150 ++++++++ src/prkit/core/model_clients/openai.py | 263 +++++++++++--- tests/prkit/core/model_clients/test_batch.py | 348 +++++++++++++++++++ 8 files changed, 1086 insertions(+), 164 deletions(-) create mode 100644 src/prkit/core/model_clients/batch_types.py create mode 100644 tests/prkit/core/model_clients/test_batch.py diff --git a/CHANGELOG.md b/CHANGELOG.md index b53603a..9959559 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,7 @@ Production releases follow semantic versioning. TestPyPI validation builds use P ### Added +- **Batch API support across OpenAI, Anthropic, and Gemini** — `BaseModelClient` gains a synchronous batch job lifecycle (`submit_batch` → `poll_batch` → `retrieve_batch_results`) plus a free-text request builder `build_batch_request(...)` that mirrors `response()` (same `input`/`instructions` handling, no structured output), complementing the existing structured `build_batch_structured_request`. New provider-agnostic types `BatchState`, `BatchStatus`, `BatchItemStatus`, and `BatchResult` (in `prkit.core.model_clients.batch_types`) normalize each provider's status enum and per-request results. Each provider's request-body construction is now shared between `response()` and the batch builders (`_build_responses_body` / `_build_messages_params`) to prevent drift. OpenAI o-family models drop `temperature` at build time. Unsupported providers raise `NotImplementedError`. Batch processing runs asynchronously at ~50% of synchronous cost. - **`OpenAIModel` custom endpoint support** — new keyword-only constructor params `base_url`, `api_key`, and `api_key_env` allow routing to any proxy or gateway that implements the OpenAI Responses API (`POST /v1/responses`) with an explicit key or key from a named environment variable. Backward-compatible: omitting all three preserves existing `OPENAI_API_KEY` + default endpoint behaviour. - **`OllamaModel` explicit auth params** — new keyword-only constructor params `api_key` and `api_key_env` forward a `Bearer` token as the `Authorization` header to `ollama.Client`, providing API-key parity with other providers. Works for cloud endpoints (e.g. `base_url="https://ollama.com"`). - **Remote-safe Ollama preflight** — when `base_url` or `OLLAMA_HOST` points to a non-local host, a failed startup connectivity check now emits a warning instead of raising `ConnectionError`; precise errors surface at `chat()` call time. diff --git a/src/prkit/core/model_clients/__init__.py b/src/prkit/core/model_clients/__init__.py index a3f93c4..bde98b8 100644 --- a/src/prkit/core/model_clients/__init__.py +++ b/src/prkit/core/model_clients/__init__.py @@ -9,6 +9,13 @@ from typing import Any from .base import DEFAULT_INSTRUCTIONS, BaseModelClient +from .batch_types import ( + TERMINAL_STATES, + BatchItemStatus, + BatchResult, + BatchState, + BatchStatus, +) from .factory import ProviderRule, create_model_client, register_model_client from .modes import PhysicsOutputMode from .prompts import build_plain_question_prompt, format_problem_context @@ -17,9 +24,14 @@ __all__ = [ "BaseModelClient", + "BatchItemStatus", + "BatchResult", + "BatchState", + "BatchStatus", "DEFAULT_INSTRUCTIONS", "PhysicsOutputMode", "ProviderRule", + "TERMINAL_STATES", "build_plain_question_prompt", "create_model_client", "create_llm_client", diff --git a/src/prkit/core/model_clients/anthropic.py b/src/prkit/core/model_clients/anthropic.py index f90715d..4cc7cce 100644 --- a/src/prkit/core/model_clients/anthropic.py +++ b/src/prkit/core/model_clients/anthropic.py @@ -4,12 +4,13 @@ import logging import os import re -from collections.abc import Callable +from collections.abc import Callable, Iterator, Sequence from typing import Any from pydantic import BaseModel from .base import BaseModelClient +from .batch_types import BatchItemStatus, BatchResult, BatchState, BatchStatus from .structured_output import ( StructuredOutputPlan, StructuredOutputPolicy, @@ -38,6 +39,21 @@ ANTHROPIC_OPTIONAL_PARAMETER_LIMIT = 24 ANTHROPIC_UNION_PARAMETER_LIMIT = 16 +# Map the Anthropic batch ``processing_status`` onto a provider-agnostic ``BatchState``. +_ANTHROPIC_BATCH_STATE_MAP = { + "in_progress": BatchState.IN_PROGRESS, + "canceling": BatchState.IN_PROGRESS, + "ended": BatchState.COMPLETED, +} + +# Map a per-request batch ``result.type`` onto a provider-agnostic ``BatchItemStatus``. +_ANTHROPIC_ITEM_STATUS_MAP = { + "succeeded": BatchItemStatus.SUCCEEDED, + "errored": BatchItemStatus.ERRORED, + "canceled": BatchItemStatus.CANCELED, + "expired": BatchItemStatus.EXPIRED, +} + def _detect_image_media_type(image_path: str) -> str: """Detect media type for image file path.""" @@ -219,69 +235,18 @@ def response( FileNotFoundError: If any image_path is a file path that doesn't exist ValueError: If a data URL is malformed """ - normalized_response_format = None - if response_format is not None: - normalized_response_format = normalize_response_format(response_format) - - content: list[dict[str, Any]] = [{"type": "text", "text": input}] - if image_paths: - for image_path in image_paths: - if image_path.startswith("data:"): - parsed = _parse_data_url(image_path) - content.append( - { - "type": "image", - "source": { - "type": "base64", - "media_type": parsed["media_type"], - "data": parsed["data"], - }, - } - ) - elif image_path.startswith("http://") or image_path.startswith( - "https://" - ): - self.logger.warning( - "Anthropic image URL inputs are not enabled in this client yet. " - f"Ignoring URL image input: {image_path}" - ) - else: - if not os.path.exists(image_path): - raise FileNotFoundError(f"Image file not found: {image_path}") - - content.append( - { - "type": "image", - "source": { - "type": "base64", - "media_type": _detect_image_media_type(image_path), - "data": encode_image_to_base64(image_path), - }, - } - ) - - request_params: dict[str, Any] = { - "model": self.model, - "messages": [{"role": "user", "content": content}], - "max_tokens": max_output_tokens, - } - instr = self._resolve_instructions(instructions) - if instr: - request_params["system"] = instr - if normalized_response_format is not None: - request_params["output_config"] = { - "format": { - "type": "json_schema", - "name": normalized_response_format["name"], - "schema": normalized_response_format["schema"], - } - } - if kwargs: - request_params.update(kwargs) + params = self._build_messages_params( + input=input, + instructions=self._resolve_instructions(instructions), + image_paths=image_paths, + max_output_tokens=max_output_tokens, + response_format=response_format, + extra=kwargs, + ) - response = self.client.messages.create(**request_params) + response = self.client.messages.create(**params) - if normalized_response_format is not None: + if response_format is not None: text = _extract_text_json(response.content) self.logger.info(f"Response: {text}") return text @@ -295,6 +260,142 @@ def response( self.logger.info(f"Response: {text}") return text + def _build_content_blocks( + self, input: str, image_paths: Sequence[str] | None + ) -> list[dict[str, Any]]: + """Build Anthropic user-message content blocks (text + base64 images).""" + content: list[dict[str, Any]] = [{"type": "text", "text": input}] + if not image_paths: + return content + for image_path in image_paths: + if image_path.startswith("data:"): + parsed = _parse_data_url(image_path) + content.append( + { + "type": "image", + "source": { + "type": "base64", + "media_type": parsed["media_type"], + "data": parsed["data"], + }, + } + ) + elif image_path.startswith("http://") or image_path.startswith("https://"): + self.logger.warning( + "Anthropic image URL inputs are not enabled in this client yet. " + f"Ignoring URL image input: {image_path}" + ) + else: + if not os.path.exists(image_path): + raise FileNotFoundError(f"Image file not found: {image_path}") + content.append( + { + "type": "image", + "source": { + "type": "base64", + "media_type": _detect_image_media_type(image_path), + "data": encode_image_to_base64(image_path), + }, + } + ) + return content + + def _build_messages_params( + self, + *, + input: str, + instructions: str | None, + image_paths: Sequence[str] | None, + max_output_tokens: int, + response_format: dict[str, Any] | type | None = None, + extra: dict[str, Any] | None = None, + ) -> dict[str, Any]: + """Build the Messages API params shared by ``response()`` and the batch builders. + + *instructions* must already be resolved (an empty string omits ``system``). + *extra* carries any remaining params (e.g. ``temperature``) and is merged last. + """ + params: dict[str, Any] = { + "model": self.model, + "messages": [ + { + "role": "user", + "content": self._build_content_blocks(input, image_paths), + } + ], + "max_tokens": max_output_tokens, + } + if instructions: + params["system"] = instructions + if response_format is not None: + normalized = normalize_response_format(response_format) + params["output_config"] = { + "format": { + "type": "json_schema", + "name": normalized["name"], + "schema": normalized["schema"], + } + } + if extra: + params.update(extra) + return params + + def _build_batch_request( + self, + *, + request_id: str, + input: str, + instructions: str | None, + image_paths: tuple[str, ...], + max_output_tokens: int | None, + temperature: float | None, + **kwargs: Any, + ) -> dict[str, Any]: + """Build a free-text Messages batch request ({custom_id, params}).""" + extra: dict[str, Any] = dict(kwargs) + if temperature is not None: + extra["temperature"] = temperature + params = self._build_messages_params( + input=input, + instructions=instructions, + image_paths=image_paths, + max_output_tokens=max_output_tokens if max_output_tokens is not None else 1024, + response_format=None, + extra=extra, + ) + return {"custom_id": request_id, "params": params} + + def submit_batch( + self, + requests: Sequence[dict[str, Any]], + *, + metadata: dict[str, str] | None = None, + ) -> str: + del metadata # Anthropic batch creation takes no metadata argument. + batch = self.client.messages.batches.create(requests=list(requests)) + return str(batch.id) + + def poll_batch(self, batch_id: str) -> BatchStatus: + batch = self.client.messages.batches.retrieve(batch_id) + raw_status = str(getattr(batch, "processing_status", "") or "") + counts: dict[str, int] = {} + request_counts = getattr(batch, "request_counts", None) + if request_counts is not None: + for key in ("processing", "succeeded", "errored", "canceled", "expired"): + counts[key] = int(getattr(request_counts, key, 0) or 0) + return BatchStatus( + batch_id=batch_id, + state=_ANTHROPIC_BATCH_STATE_MAP.get(raw_status, BatchState.UNKNOWN), + provider=self._provider_name(), + raw_status=raw_status, + counts=counts, + output_ref=getattr(batch, "results_url", None), + ) + + def retrieve_batch_results(self, batch_id: str) -> Iterator[BatchResult]: + for entry in self.client.messages.batches.results(batch_id): + yield _parse_anthropic_result_entry(entry) + def _resolve_structured_output_plan( self, spec: StructuredOutputSpec, @@ -352,47 +453,34 @@ def _build_batch_structured_request( "Anthropic batch structured requests require json_schema mode. " f"Got {plan.mode!r}." ) - - normalized = normalize_response_format(plan.response_format or {}) - content: list[dict[str, Any]] = [{"type": "text", "text": user_prompt}] - for image_path in image_paths: - if image_path.startswith("data:"): - parsed = _parse_data_url(image_path) - content.append( - { - "type": "image", - "source": { - "type": "base64", - "media_type": parsed["media_type"], - "data": parsed["data"], - }, - } - ) - else: - content.append( - { - "type": "image", - "source": { - "type": "base64", - "media_type": _detect_image_media_type(image_path), - "data": encode_image_to_base64(image_path), - }, - } - ) - - params: dict[str, Any] = { - "model": self.model, - "max_tokens": max_output_tokens or 4096, - "messages": [{"role": "user", "content": content}], - "output_config": { - "format": { - "type": "json_schema", - "schema": normalized["schema"], - } - }, - } - - return { - "custom_id": request_id, - "params": params, - } + params = self._build_messages_params( + input=user_prompt, + instructions=None, + image_paths=image_paths, + max_output_tokens=max_output_tokens if max_output_tokens is not None else 4096, + response_format=plan.response_format or {}, + ) + return {"custom_id": request_id, "params": params} + + +def _parse_anthropic_result_entry(entry: Any) -> BatchResult: + """Parse one streamed Message Batch result entry into a ``BatchResult``.""" + custom_id = str(getattr(entry, "custom_id", "") or "") + result = getattr(entry, "result", None) + result_type = str(getattr(result, "type", "") or "") + if result_type == "succeeded": + message = getattr(result, "message", None) + blocks = getattr(message, "content", None) or [] + text = "\n".join( + str(_block_attr(block, "text")) + for block in blocks + if _block_attr(block, "type") == "text" and _block_attr(block, "text") + ).strip() + return BatchResult(custom_id, BatchItemStatus.SUCCEEDED, text=text) + status = _ANTHROPIC_ITEM_STATUS_MAP.get(result_type, BatchItemStatus.ERRORED) + error = getattr(result, "error", None) + return BatchResult( + custom_id, + status, + error=str(error) if error is not None else result_type or "unknown", + ) diff --git a/src/prkit/core/model_clients/base.py b/src/prkit/core/model_clients/base.py index 1f29191..647044f 100644 --- a/src/prkit/core/model_clients/base.py +++ b/src/prkit/core/model_clients/base.py @@ -5,7 +5,7 @@ import logging import warnings from abc import ABC, abstractmethod -from collections.abc import Sequence +from collections.abc import Iterator, Sequence from typing import TYPE_CHECKING, Any, TypeVar from pydantic import BaseModel, ValidationError @@ -13,6 +13,7 @@ from ..domain import PhysicsProblem from ..logging_config import PRKitLogger from ..project_env import load_project_dotenv +from .batch_types import BatchResult, BatchStatus from .modes import PhysicsOutputMode from .prompts import build_plain_question_prompt from .structured_output import ( @@ -308,6 +309,90 @@ def parse_batch_structured_response( plan=plan, ) + # ------------------------------------------------------------------ + # Free-text batch requests + the asynchronous job lifecycle + # ------------------------------------------------------------------ + def build_batch_request( + self, + *, + request_id: str, + input: str, + instructions: str | None = None, + image_paths: Sequence[str] | None = None, + max_output_tokens: int | None = None, + temperature: float | None = None, + **kwargs: Any, + ) -> dict[str, Any]: + """Build a provider-specific FREE-TEXT batch request line (no structured output). + + Mirrors :meth:`response`: the same *input* / *instructions* handling and + no ``response_format``. Passing ``instructions=""`` suppresses the system + prompt on every provider (see :meth:`_resolve_instructions`), so the + resulting request matches a synchronous + ``response(input=..., instructions="")`` call. Use + :meth:`build_batch_structured_request` instead for schema-enforced output. + """ + return self._build_batch_request( + request_id=request_id, + input=input, + instructions=self._resolve_instructions(instructions), + image_paths=tuple(image_paths or ()), + max_output_tokens=max_output_tokens, + temperature=temperature, + **kwargs, + ) + + def _build_batch_request( + self, + *, + request_id: str, + input: str, + instructions: str | None, + image_paths: tuple[str, ...], + max_output_tokens: int | None, + temperature: float | None, + **kwargs: Any, + ) -> dict[str, Any]: + """Override per provider. Default raises ``NotImplementedError``.""" + del ( + request_id, + input, + instructions, + image_paths, + max_output_tokens, + temperature, + kwargs, + ) + raise NotImplementedError( + f"Free-text batch requests are not implemented for provider={self._provider_name()!r}." + ) + + def submit_batch( + self, + requests: Sequence[dict[str, Any]], + *, + metadata: dict[str, str] | None = None, + ) -> str: + """Submit a list of batch request dicts; return the provider batch/job id.""" + del requests, metadata + raise NotImplementedError( + f"Batch submission is not implemented for provider={self._provider_name()!r}." + ) + + def poll_batch(self, batch_id: str) -> BatchStatus: + """Return a single snapshot of the batch job's state (no sleeping/looping).""" + del batch_id + raise NotImplementedError( + f"Batch polling is not implemented for provider={self._provider_name()!r}." + ) + + def retrieve_batch_results(self, batch_id: str) -> Iterator[BatchResult]: + """Yield per-request results for a terminal batch, correlated by ``custom_id``.""" + del batch_id + raise NotImplementedError( + f"Batch result retrieval is not implemented for provider={self._provider_name()!r}." + ) + def _resolve_structured_output_plan( self, spec: StructuredOutputSpec, diff --git a/src/prkit/core/model_clients/batch_types.py b/src/prkit/core/model_clients/batch_types.py new file mode 100644 index 0000000..f95bbf0 --- /dev/null +++ b/src/prkit/core/model_clients/batch_types.py @@ -0,0 +1,89 @@ +"""Provider-agnostic types for the asynchronous Batch API lifecycle. + +Batch processing is shaped differently from a synchronous chat call: many +requests are submitted at once, the provider returns a job id, the caller polls +until the job reaches a *terminal* state, then downloads per-request results +(each of which may independently succeed or fail). These types give OpenAI, +Anthropic, and Gemini batches one common vocabulary so a consumer can drive all +three through the same interface (see ``BaseModelClient.submit_batch`` / +``poll_batch`` / ``retrieve_batch_results``). +""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +from .modes import _StrEnum + + +class BatchState(_StrEnum): + """Provider-agnostic state of a whole batch job.""" + + PENDING = "pending" # accepted / validating, not yet running + IN_PROGRESS = "in_progress" # running (incl. finalizing / cancelling) + COMPLETED = "completed" # terminal: finished, results available + FAILED = "failed" # terminal: batch-level failure + EXPIRED = "expired" # terminal: window elapsed before completion + CANCELLED = "cancelled" # terminal: cancelled by the user + UNKNOWN = "unknown" # unrecognized provider status + + +TERMINAL_STATES = frozenset( + { + BatchState.COMPLETED, + BatchState.FAILED, + BatchState.EXPIRED, + BatchState.CANCELLED, + } +) + + +class BatchItemStatus(_StrEnum): + """Per-request outcome within a batch that has finished processing.""" + + SUCCEEDED = "succeeded" + ERRORED = "errored" + EXPIRED = "expired" + CANCELED = "canceled" + + +@dataclass(frozen=True) +class BatchStatus: + """A snapshot of a batch job's state from a single poll. + + *output_ref* / *error_ref* carry whatever handle the provider needs to fetch + results later (e.g. an OpenAI output file id, a Gemini destination file + name); they may be ``None`` until the job is terminal. + """ + + batch_id: str + state: BatchState + provider: str + raw_status: str + counts: dict[str, int] = field(default_factory=dict) + output_ref: str | None = None + error_ref: str | None = None + + @property + def is_terminal(self) -> bool: + """True once the job has stopped processing (success or failure).""" + return self.state in TERMINAL_STATES + + +@dataclass(frozen=True) +class BatchResult: + """One request's result, correlated back to its input by ``custom_id``. + + On success *text* holds the model's free-text output (the same string a + synchronous ``response()`` would return). On any non-success outcome *text* + is ``None`` and *error* describes the failure. + """ + + custom_id: str + status: BatchItemStatus + text: str | None = None + error: str | None = None + + @property + def succeeded(self) -> bool: + return self.status is BatchItemStatus.SUCCEEDED diff --git a/src/prkit/core/model_clients/gemini.py b/src/prkit/core/model_clients/gemini.py index 6487efe..959e3b3 100644 --- a/src/prkit/core/model_clients/gemini.py +++ b/src/prkit/core/model_clients/gemini.py @@ -1,7 +1,9 @@ """Google Gemini API model client using the ``google-genai`` SDK.""" +import json import logging import os +from collections.abc import Iterator, Sequence from typing import Any import PIL.Image @@ -10,6 +12,7 @@ from pydantic import BaseModel from .base import BaseModelClient +from .batch_types import BatchItemStatus, BatchResult, BatchState, BatchStatus from .structured_output import ( StructuredOutputPlan, StructuredOutputPolicy, @@ -19,6 +22,17 @@ ) from .utils import detect_image_mime_type, encode_image_to_base64 +# Map the Gemini batch job ``state.name`` onto a provider-agnostic ``BatchState``. +_GEMINI_BATCH_STATE_MAP = { + "JOB_STATE_PENDING": BatchState.PENDING, + "JOB_STATE_QUEUED": BatchState.PENDING, + "JOB_STATE_RUNNING": BatchState.IN_PROGRESS, + "JOB_STATE_SUCCEEDED": BatchState.COMPLETED, + "JOB_STATE_FAILED": BatchState.FAILED, + "JOB_STATE_CANCELLED": BatchState.CANCELLED, + "JOB_STATE_EXPIRED": BatchState.EXPIRED, +} + class GeminiModel(BaseModelClient): """Google Gemini API client implementation.""" @@ -195,6 +209,142 @@ def _build_batch_structured_request( }, } + def _build_batch_request( + self, + *, + request_id: str, + input: str, + instructions: str | None, + image_paths: tuple[str, ...], + max_output_tokens: int | None, + temperature: float | None, + **kwargs: Any, + ) -> dict[str, Any]: + """Build a free-text Gemini batch request line ({key, request}). + + Mirrors a synchronous ``response(input=..., instructions=...)`` call: the + folded prompt becomes the user text, *instructions* (when truthy) becomes + the request-level ``system_instruction``, and no ``response_json_schema`` + is set (free-text output). + """ + parts: list[dict[str, Any]] = [{"text": input}] + for image_path in image_paths: + parts.append( + { + "inline_data": { + "mime_type": _guess_mime_type(image_path), + "data": _encode_image_file_base64(image_path), + } + } + ) + + generation_config: dict[str, Any] = {} + if max_output_tokens is not None: + generation_config["max_output_tokens"] = max_output_tokens + if temperature is not None: + generation_config["temperature"] = temperature + generation_config.update(kwargs) + + request: dict[str, Any] = {"contents": [{"parts": parts, "role": "user"}]} + if generation_config: + request["generation_config"] = generation_config + if instructions: + request["system_instruction"] = {"parts": [{"text": instructions}]} + return {"key": request_id, "request": request} + + def submit_batch( + self, + requests: Sequence[dict[str, Any]], + *, + metadata: dict[str, str] | None = None, + ) -> str: + """Create a per-model inline batch job; return its job ``name``.""" + create_kwargs: dict[str, Any] = {"model": self.model, "src": list(requests)} + if metadata: + create_kwargs["config"] = { + "display_name": metadata.get("display_name", "prkit-batch") + } + job = self.genai_client.batches.create(**create_kwargs) + return str(job.name) + + def poll_batch(self, batch_id: str) -> BatchStatus: + job = self.genai_client.batches.get(name=batch_id) + state = getattr(job, "state", None) + raw_status = getattr(state, "name", None) or str(state or "") + dest = getattr(job, "dest", None) + output_ref = getattr(dest, "file_name", None) if dest is not None else None + return BatchStatus( + batch_id=batch_id, + state=_GEMINI_BATCH_STATE_MAP.get(raw_status, BatchState.UNKNOWN), + provider=self._provider_name(), + raw_status=raw_status, + output_ref=output_ref, + ) + + def retrieve_batch_results(self, batch_id: str) -> Iterator[BatchResult]: + job = self.genai_client.batches.get(name=batch_id) + dest = getattr(job, "dest", None) + inlined = getattr(dest, "inlined_responses", None) if dest is not None else None + if inlined: + for item in inlined: + yield _parse_gemini_inline_response(item) + return + file_name = getattr(dest, "file_name", None) if dest is not None else None + if file_name: + content = self.genai_client.files.download(file=file_name) + text = ( + content.decode("utf-8") + if isinstance(content, (bytes, bytearray)) + else str(content) + ) + for line in text.splitlines(): + stripped = line.strip() + if stripped: + yield _parse_gemini_result_line(stripped) + + +def _extract_gemini_text(response: Any) -> str: + """Extract the assistant text from a Gemini response (dict from file, or SDK object).""" + if isinstance(response, dict): + for candidate in response.get("candidates") or []: + content = (candidate or {}).get("content") or {} + chunks = [ + str(part.get("text")) + for part in content.get("parts") or [] + if isinstance(part, dict) and part.get("text") + ] + if chunks: + return "".join(chunks) + return "" + text = getattr(response, "text", None) + return str(text) if text else "" + + +def _parse_gemini_inline_response(item: Any) -> BatchResult: + """Parse one ``dest.inlined_responses`` entry into a ``BatchResult``.""" + key = str(getattr(item, "key", "") or "") + error = getattr(item, "error", None) + response = getattr(item, "response", None) + if error is not None or response is None: + return BatchResult( + key, + BatchItemStatus.ERRORED, + error=str(error) if error is not None else "no response", + ) + return BatchResult(key, BatchItemStatus.SUCCEEDED, text=_extract_gemini_text(response)) + + +def _parse_gemini_result_line(line: str) -> BatchResult: + """Parse one JSONL line ({key, response|error}) from a batch output file.""" + obj = json.loads(line) + key = str(obj.get("key", "") or "") + error = obj.get("error") + response = obj.get("response") + if error is not None or response is None: + message = json.dumps(error) if error is not None else "no response" + return BatchResult(key, BatchItemStatus.ERRORED, error=message) + return BatchResult(key, BatchItemStatus.SUCCEEDED, text=_extract_gemini_text(response)) + def _extract_gemini_error_details(response: object) -> str | None: """Extract block reason and error details from Gemini response when text is empty.""" diff --git a/src/prkit/core/model_clients/openai.py b/src/prkit/core/model_clients/openai.py index cf926af..fe35986 100644 --- a/src/prkit/core/model_clients/openai.py +++ b/src/prkit/core/model_clients/openai.py @@ -9,8 +9,11 @@ - o-family (o3, o4, o4-mini, etc. - models starting with 'o' followed by number) """ +import io +import json import logging import os +from collections.abc import Iterator, Sequence from typing import Any from openai import OpenAI @@ -18,6 +21,7 @@ from ..project_env import ensure_openai_api_key from .base import BaseModelClient +from .batch_types import BatchItemStatus, BatchResult, BatchState, BatchStatus from .structured_output import ( StructuredOutputPlan, StructuredOutputPolicy, @@ -27,6 +31,18 @@ ) from .utils import prepare_image_url_from_path +# Map the OpenAI batch ``status`` string onto a provider-agnostic ``BatchState``. +_OPENAI_BATCH_STATE_MAP = { + "validating": BatchState.PENDING, + "in_progress": BatchState.IN_PROGRESS, + "finalizing": BatchState.IN_PROGRESS, + "completed": BatchState.COMPLETED, + "failed": BatchState.FAILED, + "expired": BatchState.EXPIRED, + "cancelling": BatchState.IN_PROGRESS, + "cancelled": BatchState.CANCELLED, +} + def _ensure_additional_properties_false(schema: dict[str, Any]) -> dict[str, Any]: """Recursively set ``additionalProperties: false`` on all object nodes.""" @@ -274,10 +290,44 @@ def response( FileNotFoundError: If any image_path is a file path that doesn't exist IOError: If there's an error reading any image file """ - # Build request parameters - request_params: dict[str, Any] = {"model": self.model} + max_output_tokens = kwargs.pop("max_output_tokens", None) + request_params = self._build_responses_body( + input=input, + instructions=self._resolve_instructions(instructions), + image_paths=image_paths, + max_output_tokens=max_output_tokens, + response_format=response_format, + extra=kwargs, + ) + + response = self.client.responses.create(**request_params) + text = str(response.output_text) + self.logger.info(f"Response: {text}") + return text + + def _omit_temperature(self) -> bool: + """o-family reasoning models reject an explicit ``temperature`` parameter.""" + return self.is_o_family + + def _build_responses_body( + self, + *, + input: str, + instructions: str | None, + image_paths: Sequence[str] | None, + max_output_tokens: int | None, + response_format: dict[str, Any] | type | None = None, + extra: dict[str, Any] | None = None, + ) -> dict[str, Any]: + """Build the Responses API body shared by ``response()`` and the batch builders. + + *instructions* must already be resolved (see ``_resolve_instructions``); an + empty string is omitted. *extra* carries any remaining request params + (e.g. ``temperature``) and is merged last, mirroring ``response()``'s + historical ``request_params.update(kwargs)`` behavior. + """ + body: dict[str, Any] = {"model": self.model} - # Add structured output if requested if response_format is not None: normalized = normalize_response_format(response_format) strict_schema = ensure_openai_strict_json_schema(normalized["schema"]) @@ -285,41 +335,111 @@ def response( "type": "json_schema", "name": normalized["name"], "schema": strict_schema, - "strict": normalized["strict"], + "strict": normalized.get("strict", True), } - request_params["text"] = {"format": text_format} if normalized.get("description") is not None: text_format["description"] = normalized["description"] + body["text"] = {"format": text_format} - # Use role/content format for all models content: list[dict[str, Any]] = [{"type": "input_text", "text": input}] - if image_paths: for image_path in image_paths: - image_url = prepare_image_url_from_image_path(image_path) - content.append({"type": "input_image", "image_url": image_url}) - - request_params["input"] = [{"role": "user", "content": content}] - - instr = self._resolve_instructions(instructions) - if instr: - request_params["instructions"] = instr + content.append( + { + "type": "input_image", + "image_url": prepare_image_url_from_image_path(image_path), + } + ) + body["input"] = [{"role": "user", "content": content}] - # Add reasoning parameter for o-family models + if instructions: + body["instructions"] = instructions if self.is_o_family: - request_params["reasoning"] = {"effort": "medium"} - - max_output_tokens = kwargs.pop("max_output_tokens", None) + body["reasoning"] = {"effort": "medium"} if max_output_tokens is not None: - request_params["max_output_tokens"] = max_output_tokens + body["max_output_tokens"] = max_output_tokens + if extra: + body.update(extra) + return body + + def _build_batch_request( + self, + *, + request_id: str, + input: str, + instructions: str | None, + image_paths: tuple[str, ...], + max_output_tokens: int | None, + temperature: float | None, + **kwargs: Any, + ) -> dict[str, Any]: + """Build a free-text Responses batch line (one ``input.jsonl`` row).""" + extra: dict[str, Any] = dict(kwargs) + if temperature is not None and not self._omit_temperature(): + extra["temperature"] = temperature + body = self._build_responses_body( + input=input, + instructions=instructions, + image_paths=image_paths, + max_output_tokens=max_output_tokens, + response_format=None, + extra=extra, + ) + return { + "custom_id": request_id, + "method": "POST", + "url": "/v1/responses", + "body": body, + } - if kwargs: - request_params.update(kwargs) + def submit_batch( + self, + requests: Sequence[dict[str, Any]], + *, + metadata: dict[str, str] | None = None, + ) -> str: + """Upload the requests as a JSONL file and create a 24h batch job.""" + payload = "\n".join(json.dumps(request) for request in requests) + upload = self.client.files.create( + file=("batch_requests.jsonl", io.BytesIO(payload.encode("utf-8"))), + purpose="batch", + ) + batch = self.client.batches.create( + input_file_id=upload.id, + endpoint="/v1/responses", + completion_window="24h", + metadata=metadata or {}, + ) + return str(batch.id) + + def poll_batch(self, batch_id: str) -> BatchStatus: + batch = self.client.batches.retrieve(batch_id) + raw_status = str(getattr(batch, "status", "") or "") + counts: dict[str, int] = {} + request_counts = getattr(batch, "request_counts", None) + if request_counts is not None: + for key in ("total", "completed", "failed"): + counts[key] = int(getattr(request_counts, key, 0) or 0) + return BatchStatus( + batch_id=batch_id, + state=_OPENAI_BATCH_STATE_MAP.get(raw_status, BatchState.UNKNOWN), + provider=self._provider_name(), + raw_status=raw_status, + counts=counts, + output_ref=getattr(batch, "output_file_id", None), + error_ref=getattr(batch, "error_file_id", None), + ) - response = self.client.responses.create(**request_params) - text = str(response.output_text) - self.logger.info(f"Response: {text}") - return text + def retrieve_batch_results(self, batch_id: str) -> Iterator[BatchResult]: + batch = self.client.batches.retrieve(batch_id) + output_file_id = getattr(batch, "output_file_id", None) + if output_file_id: + for line in _iter_jsonl_lines(self.client.files.content(output_file_id)): + yield _parse_openai_result_line(line) + error_file_id = getattr(batch, "error_file_id", None) + if error_file_id: + for line in _iter_jsonl_lines(self.client.files.content(error_file_id)): + yield _parse_openai_result_line(line, force_error=True) def _resolve_structured_output_plan( self, @@ -393,40 +513,69 @@ def _build_batch_structured_request( raise ValueError( f"OpenAI batch structured requests require json_schema mode. Got {plan.mode!r}." ) - - normalized = normalize_response_format(plan.response_format or {}) - strict_schema = ensure_openai_strict_json_schema(normalized["schema"]) - text_format: dict[str, Any] = { - "type": "json_schema", - "name": normalized["name"], - "schema": strict_schema, - "strict": normalized.get("strict", True), - } - if normalized.get("description") is not None: - text_format["description"] = normalized["description"] - - content: list[dict[str, Any]] = [{"type": "input_text", "text": user_prompt}] - for image_path in image_paths: - content.append( - { - "type": "input_image", - "image_url": prepare_image_url_from_image_path(image_path), - } - ) - - body: dict[str, Any] = { - "model": self.model, - "input": [{"role": "user", "content": content}], - "text": {"format": text_format}, - } - if self.is_o_family: - body["reasoning"] = {"effort": "medium"} - if max_output_tokens is not None: - body["max_output_tokens"] = max_output_tokens - + body = self._build_responses_body( + input=user_prompt, + instructions=None, + image_paths=image_paths, + max_output_tokens=max_output_tokens, + response_format=plan.response_format or {}, + ) return { "custom_id": request_id, "method": "POST", "url": "/v1/responses", "body": body, } + + +def _iter_jsonl_lines(content: Any) -> Iterator[str]: + """Yield non-empty lines from an OpenAI ``files.content`` payload. + + ``files.content`` returns an ``HttpxBinaryResponseContent``; prefer its + ``.text``, falling back to ``.read()`` / raw bytes for stubbed clients. + """ + text = getattr(content, "text", None) + if text is None: + raw = content.read() if hasattr(content, "read") else content + text = raw.decode("utf-8") if isinstance(raw, (bytes, bytearray)) else str(raw) + for line in text.splitlines(): + stripped = line.strip() + if stripped: + yield stripped + + +def _extract_output_text_from_responses_body(body: Any) -> str: + """Pull the assistant text out of a Responses API object body (batch output line).""" + if not isinstance(body, dict): + return "" + flat = body.get("output_text") + if isinstance(flat, str) and flat: + return flat + chunks: list[str] = [] + for item in body.get("output") or []: + if not isinstance(item, dict) or item.get("type") != "message": + continue + for part in item.get("content") or []: + if isinstance(part, dict) and part.get("type") == "output_text": + piece = part.get("text") + if piece: + chunks.append(str(piece)) + return "".join(chunks) + + +def _parse_openai_result_line(line: str, *, force_error: bool = False) -> BatchResult: + """Parse one JSONL line from a batch output/error file into a ``BatchResult``.""" + obj = json.loads(line) + custom_id = str(obj.get("custom_id", "")) + error = obj.get("error") + response = obj.get("response") if isinstance(obj.get("response"), dict) else None + status_code = response.get("status_code") if response else None + if force_error or error is not None or (status_code is not None and status_code != 200): + message = json.dumps(error) if error is not None else f"status_code={status_code}" + return BatchResult(custom_id, BatchItemStatus.ERRORED, error=message) + body = response.get("body") if response else None + return BatchResult( + custom_id, + BatchItemStatus.SUCCEEDED, + text=_extract_output_text_from_responses_body(body), + ) diff --git a/tests/prkit/core/model_clients/test_batch.py b/tests/prkit/core/model_clients/test_batch.py new file mode 100644 index 0000000..e81251a --- /dev/null +++ b/tests/prkit/core/model_clients/test_batch.py @@ -0,0 +1,348 @@ +"""Tests for free-text batch request builders and the batch job lifecycle. + +Provider SDKs are faked with ``MagicMock`` so these run fully offline. The +lifecycle tests assert the status-enum mapping and per-result text extraction +for each provider; the parity tests assert a batch request body equals the body +the synchronous ``response()`` path would build from the same inputs. +""" + +from __future__ import annotations + +import json +from unittest.mock import MagicMock, patch + +import pytest + +from prkit.core.model_clients.base import BaseModelClient +from prkit.core.model_clients.batch_types import ( + TERMINAL_STATES, + BatchItemStatus, + BatchState, + BatchStatus, +) + + +def _openai_client(model: str = "gpt-5.4-mini"): + with patch("prkit.core.model_clients.openai.OpenAI") as mock_cls: + mock_cls.return_value = MagicMock() + from prkit.core.model_clients.openai import OpenAIModel + + return OpenAIModel(model) + + +def _anthropic_client(model: str = "claude-opus-4-8"): + with patch("prkit.core.model_clients.anthropic.Anthropic") as mock_cls: + mock_cls.return_value = MagicMock() + from prkit.core.model_clients.anthropic import AnthropicModel + + return AnthropicModel(model) + + +def _gemini_client(model: str = "gemini-3.5-flash"): + with patch("prkit.core.model_clients.gemini.genai.Client") as mock_cls: + mock_cls.return_value = MagicMock() + from prkit.core.model_clients.gemini import GeminiModel + + return GeminiModel(model) + + +class TestBatchTypes: + def test_terminal_states(self): + assert BatchState.COMPLETED in TERMINAL_STATES + assert BatchState.FAILED in TERMINAL_STATES + assert BatchState.IN_PROGRESS not in TERMINAL_STATES + + def test_status_is_terminal(self): + assert BatchStatus("b", BatchState.COMPLETED, "openai", "completed").is_terminal + assert not BatchStatus("b", BatchState.PENDING, "openai", "validating").is_terminal + + +class TestBaseDefaults: + def test_batch_methods_raise_not_implemented(self): + class Dummy(BaseModelClient): + def response(self, input, image_paths=None, response_format=None, *, instructions=None, **kwargs): + return "" + + d = Dummy("dummy-model") + with pytest.raises(NotImplementedError): + d.build_batch_request(request_id="r", input="x") + with pytest.raises(NotImplementedError): + d.submit_batch([]) + with pytest.raises(NotImplementedError): + d.poll_batch("b") + with pytest.raises(NotImplementedError): + list(d.retrieve_batch_results("b")) + + +class TestFreeTextBuilders: + def test_openai_shape(self): + client = _openai_client() + req = client.build_batch_request( + request_id="req_1", input="hello", instructions="", max_output_tokens=100, temperature=0.0 + ) + assert req["custom_id"] == "req_1" + assert req["method"] == "POST" + assert req["url"] == "/v1/responses" + body = req["body"] + assert body["model"] == "gpt-5.4-mini" + assert body["input"][0]["content"][0]["text"] == "hello" + assert "text" not in body # no structured response_format + assert "instructions" not in body # instructions="" suppressed (parity hinge) + assert body["max_output_tokens"] == 100 + assert body["temperature"] == 0.0 + + def test_openai_o_family_drops_temperature_sets_reasoning(self): + client = _openai_client("o3") + body = client.build_batch_request( + request_id="r", input="x", instructions="", max_output_tokens=50, temperature=0.7 + )["body"] + assert "temperature" not in body + assert body["reasoning"] == {"effort": "medium"} + + def test_anthropic_shape(self): + client = _anthropic_client() + params = client.build_batch_request( + request_id="req_2", input="hi", instructions="", max_output_tokens=64, temperature=0.2 + )["params"] + assert params["model"] == "claude-opus-4-8" + assert params["max_tokens"] == 64 + assert params["messages"][0]["content"][0]["text"] == "hi" + assert "output_config" not in params + assert "system" not in params + assert params["temperature"] == 0.2 + + def test_anthropic_defaults_max_tokens_when_none(self): + client = _anthropic_client() + params = client.build_batch_request(request_id="r", input="x", instructions="")["params"] + assert params["max_tokens"] == 1024 + + def test_gemini_shape(self): + client = _gemini_client() + req = client.build_batch_request( + request_id="req_3", input="q", instructions="", max_output_tokens=32, temperature=0.5 + ) + assert req["key"] == "req_3" + request = req["request"] + assert request["contents"][0]["parts"][0]["text"] == "q" + assert request["generation_config"]["max_output_tokens"] == 32 + assert request["generation_config"]["temperature"] == 0.5 + assert "response_json_schema" not in request["generation_config"] + assert "system_instruction" not in request + + def test_instructions_included_when_provided(self): + o = _openai_client().build_batch_request( + request_id="a", input="x", instructions="SYS", max_output_tokens=10 + ) + assert o["body"]["instructions"] == "SYS" + a = _anthropic_client().build_batch_request( + request_id="a", input="x", instructions="SYS", max_output_tokens=10 + ) + assert a["params"]["system"] == "SYS" + g = _gemini_client().build_batch_request( + request_id="a", input="x", instructions="SYS", max_output_tokens=10 + ) + assert g["request"]["system_instruction"]["parts"][0]["text"] == "SYS" + + +class TestOpenAILifecycle: + def test_submit_uploads_then_creates(self): + client = _openai_client() + client.client.files.create.return_value = MagicMock(id="file_123") + client.client.batches.create.return_value = MagicMock(id="batch_abc") + bid = client.submit_batch( + [{"custom_id": "r1", "method": "POST", "url": "/v1/responses", "body": {}}] + ) + assert bid == "batch_abc" + _, kwargs = client.client.batches.create.call_args + assert kwargs["input_file_id"] == "file_123" + assert kwargs["endpoint"] == "/v1/responses" + assert kwargs["completion_window"] == "24h" + + def test_poll_maps_status_and_counts(self): + client = _openai_client() + client.client.batches.retrieve.return_value = MagicMock( + status="in_progress", + output_file_id=None, + error_file_id=None, + request_counts=MagicMock(total=2, completed=1, failed=0), + ) + status = client.poll_batch("batch_abc") + assert status.state == BatchState.IN_PROGRESS + assert status.counts["completed"] == 1 + assert not status.is_terminal + + def test_retrieve_extracts_output_text(self): + client = _openai_client() + line = json.dumps( + { + "custom_id": "r1", + "response": { + "status_code": 200, + "body": { + "output": [ + { + "type": "message", + "content": [{"type": "output_text", "text": "ANSWER"}], + } + ] + }, + }, + "error": None, + } + ) + client.client.batches.retrieve.return_value = MagicMock( + output_file_id="out_1", error_file_id=None + ) + client.client.files.content.return_value = MagicMock(text=line) + results = list(client.retrieve_batch_results("batch_abc")) + assert len(results) == 1 + assert results[0].custom_id == "r1" + assert results[0].status == BatchItemStatus.SUCCEEDED + assert results[0].text == "ANSWER" + + def test_retrieve_flags_http_error(self): + client = _openai_client() + line = json.dumps( + {"custom_id": "r2", "response": {"status_code": 500, "body": {}}, "error": None} + ) + client.client.batches.retrieve.return_value = MagicMock( + output_file_id="out_1", error_file_id=None + ) + client.client.files.content.return_value = MagicMock(text=line) + results = list(client.retrieve_batch_results("batch_abc")) + assert results[0].status == BatchItemStatus.ERRORED + + +class TestAnthropicLifecycle: + def test_submit(self): + client = _anthropic_client() + client.client.messages.batches.create.return_value = MagicMock(id="msgbatch_1") + assert client.submit_batch([{"custom_id": "r1", "params": {}}]) == "msgbatch_1" + + def test_poll(self): + client = _anthropic_client() + client.client.messages.batches.retrieve.return_value = MagicMock( + processing_status="ended", + results_url="http://results", + request_counts=MagicMock( + processing=0, succeeded=2, errored=0, canceled=0, expired=0 + ), + ) + status = client.poll_batch("msgbatch_1") + assert status.state == BatchState.COMPLETED + assert status.is_terminal + assert status.output_ref == "http://results" + assert status.counts["succeeded"] == 2 + + def test_retrieve_success(self): + client = _anthropic_client() + entry = MagicMock( + custom_id="r1", + result=MagicMock( + type="succeeded", + message=MagicMock(content=[MagicMock(type="text", text="GRADED")]), + ), + ) + client.client.messages.batches.results.return_value = iter([entry]) + results = list(client.retrieve_batch_results("msgbatch_1")) + assert results[0].custom_id == "r1" + assert results[0].text == "GRADED" + + def test_retrieve_errored(self): + client = _anthropic_client() + entry = MagicMock( + custom_id="r2", result=MagicMock(type="errored", error="boom") + ) + client.client.messages.batches.results.return_value = iter([entry]) + results = list(client.retrieve_batch_results("msgbatch_1")) + assert results[0].status == BatchItemStatus.ERRORED + + +class TestGeminiLifecycle: + def test_submit(self): + client = _gemini_client() + job = MagicMock() + job.name = "batches/xyz" + client.genai_client.batches.create.return_value = job + assert client.submit_batch([{"key": "r1", "request": {}}]) == "batches/xyz" + _, kwargs = client.genai_client.batches.create.call_args + assert kwargs["model"] == "gemini-3.5-flash" + + def test_poll(self): + client = _gemini_client() + job = MagicMock() + job.state = MagicMock() + job.state.name = "JOB_STATE_SUCCEEDED" + job.dest = MagicMock(file_name="files/out") + client.genai_client.batches.get.return_value = job + status = client.poll_batch("batches/xyz") + assert status.state == BatchState.COMPLETED + assert status.output_ref == "files/out" + + def test_retrieve_inline(self): + client = _gemini_client() + item = MagicMock() + item.key = "r1" + item.error = None + item.response = MagicMock(text="GENTEXT") + job = MagicMock() + job.dest = MagicMock() + job.dest.inlined_responses = [item] + client.genai_client.batches.get.return_value = job + results = list(client.retrieve_batch_results("batches/xyz")) + assert results[0].custom_id == "r1" + assert results[0].text == "GENTEXT" + + def test_retrieve_file(self): + client = _gemini_client() + job = MagicMock() + job.dest = MagicMock() + job.dest.inlined_responses = None + job.dest.file_name = "files/out" + client.genai_client.batches.get.return_value = job + line = json.dumps( + { + "key": "r2", + "response": {"candidates": [{"content": {"parts": [{"text": "FILETEXT"}]}}]}, + } + ) + client.genai_client.files.download.return_value = line.encode("utf-8") + results = list(client.retrieve_batch_results("batches/xyz")) + assert results[0].custom_id == "r2" + assert results[0].text == "FILETEXT" + + +class TestSyncBatchParity: + """A batch request body must equal the body sync ``response()`` would build.""" + + def test_openai_body_parity(self): + client = _openai_client() + instr = client._resolve_instructions("") + sync_body = client._build_responses_body( + input="P", + instructions=instr, + image_paths=None, + max_output_tokens=123, + response_format=None, + extra={"temperature": 0.0}, + ) + batch_body = client.build_batch_request( + request_id="r", input="P", instructions="", max_output_tokens=123, temperature=0.0 + )["body"] + assert batch_body == sync_body + + def test_anthropic_params_parity(self): + client = _anthropic_client() + instr = client._resolve_instructions("") + sync_params = client._build_messages_params( + input="P", + instructions=instr, + image_paths=None, + max_output_tokens=64, + response_format=None, + extra={"temperature": 0.2}, + ) + batch_params = client.build_batch_request( + request_id="r", input="P", instructions="", max_output_tokens=64, temperature=0.2 + )["params"] + assert batch_params == sync_params From 9e719e1a9b4a4f8eff52936a6601a0ff0137764c Mon Sep 17 00:00:00 2001 From: Yinghuan Zhang Date: Mon, 15 Jun 2026 12:23:58 -0400 Subject: [PATCH 4/8] Fix Gemini batch result correlation via file-based submission Submit Gemini batches as an uploaded keyed JSONL file (File API, src=) instead of inline requests, so results return as documented keyed JSONL and correlate to each request by key. Verified against google-genai: inline responses carry no key/metadata, so inline correlation is impossible -- the inline submit/retrieve paths are removed. - submit_batch: serialize requests to JSONL, upload, submit src= - retrieve_batch_results: file-only keyed-JSONL path - remove dead _parse_gemini_inline_response; trim _extract_gemini_text to dicts - tests: assert keyed JSONL upload + src is the file name; drop inline test - CHANGELOG: note Gemini file-based submission Co-Authored-By: Claude Opus 4.8 (1M context) --- CHANGELOG.md | 2 +- src/prkit/core/model_clients/gemini.py | 108 ++++++++--------- tests/prkit/core/model_clients/test_batch.py | 115 ++++++++++++++----- 3 files changed, 140 insertions(+), 85 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 9959559..85e1c1c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,7 +9,7 @@ Production releases follow semantic versioning. TestPyPI validation builds use P ### Added -- **Batch API support across OpenAI, Anthropic, and Gemini** — `BaseModelClient` gains a synchronous batch job lifecycle (`submit_batch` → `poll_batch` → `retrieve_batch_results`) plus a free-text request builder `build_batch_request(...)` that mirrors `response()` (same `input`/`instructions` handling, no structured output), complementing the existing structured `build_batch_structured_request`. New provider-agnostic types `BatchState`, `BatchStatus`, `BatchItemStatus`, and `BatchResult` (in `prkit.core.model_clients.batch_types`) normalize each provider's status enum and per-request results. Each provider's request-body construction is now shared between `response()` and the batch builders (`_build_responses_body` / `_build_messages_params`) to prevent drift. OpenAI o-family models drop `temperature` at build time. Unsupported providers raise `NotImplementedError`. Batch processing runs asynchronously at ~50% of synchronous cost. +- **Batch API support across OpenAI, Anthropic, and Gemini** — `BaseModelClient` gains a synchronous batch job lifecycle (`submit_batch` → `poll_batch` → `retrieve_batch_results`) plus a free-text request builder `build_batch_request(...)` that mirrors `response()` (same `input`/`instructions` handling, no structured output), complementing the existing structured `build_batch_structured_request`. New provider-agnostic types `BatchState`, `BatchStatus`, `BatchItemStatus`, and `BatchResult` (in `prkit.core.model_clients.batch_types`) normalize each provider's status enum and per-request results. Each provider's request-body construction is now shared between `response()` and the batch builders (`_build_responses_body` / `_build_messages_params`) to prevent drift. OpenAI o-family models drop `temperature` at build time. Unsupported providers raise `NotImplementedError`. Batch processing runs asynchronously at ~50% of synchronous cost. Gemini batches are submitted as an uploaded keyed JSONL file (via the File API, `src=`) rather than as inline requests, so results come back as documented keyed JSONL (`{"key": ..., "response": {...}}`) and correlate reliably to each request — inline responses carry no per-request key and cannot be correlated. - **`OpenAIModel` custom endpoint support** — new keyword-only constructor params `base_url`, `api_key`, and `api_key_env` allow routing to any proxy or gateway that implements the OpenAI Responses API (`POST /v1/responses`) with an explicit key or key from a named environment variable. Backward-compatible: omitting all three preserves existing `OPENAI_API_KEY` + default endpoint behaviour. - **`OllamaModel` explicit auth params** — new keyword-only constructor params `api_key` and `api_key_env` forward a `Bearer` token as the `Authorization` header to `ollama.Client`, providing API-key parity with other providers. Works for cloud endpoints (e.g. `base_url="https://ollama.com"`). - **Remote-safe Ollama preflight** — when `base_url` or `OLLAMA_HOST` points to a non-local host, a failed startup connectivity check now emits a warning instead of raising `ConnectionError`; precise errors surface at `chat()` call time. diff --git a/src/prkit/core/model_clients/gemini.py b/src/prkit/core/model_clients/gemini.py index 959e3b3..d7f8b13 100644 --- a/src/prkit/core/model_clients/gemini.py +++ b/src/prkit/core/model_clients/gemini.py @@ -1,5 +1,6 @@ """Google Gemini API model client using the ``google-genai`` SDK.""" +import io import json import logging import os @@ -258,13 +259,25 @@ def submit_batch( *, metadata: dict[str, str] | None = None, ) -> str: - """Create a per-model inline batch job; return its job ``name``.""" - create_kwargs: dict[str, Any] = {"model": self.model, "src": list(requests)} - if metadata: - create_kwargs["config"] = { - "display_name": metadata.get("display_name", "prkit-batch") - } - job = self.genai_client.batches.create(**create_kwargs) + """Upload the requests as keyed JSONL and create a file-based batch job. + + Each request is a ``{"key": ..., "request": {...}}`` line — the documented + Gemini batch file format. Submitting via an uploaded file (rather than + ``src=[inline requests]``) is what makes results come back as keyed JSONL + (``{"key": ..., "response": {...}}``); inline responses carry no key and + cannot be correlated back to their request. + """ + payload = "\n".join(json.dumps(request) for request in requests) + display_name = (metadata or {}).get("display_name", "prkit-batch") + uploaded = self.genai_client.files.upload( + file=io.BytesIO(payload.encode("utf-8")), + config=types.UploadFileConfig(mime_type="jsonl", display_name=display_name), + ) + job = self.genai_client.batches.create( + model=self.model, + src=uploaded.name, + config={"display_name": display_name}, + ) return str(job.name) def poll_batch(self, batch_id: str) -> BatchStatus: @@ -282,56 +295,41 @@ def poll_batch(self, batch_id: str) -> BatchStatus: ) def retrieve_batch_results(self, batch_id: str) -> Iterator[BatchResult]: + """Yield per-request results from the batch's output JSONL file. + + Submission always uses the file-based path (see ``submit_batch``), so a + terminal job exposes its results as ``dest.file_name`` — a keyed JSONL + file whose lines correlate back to each request by ``key``. + """ job = self.genai_client.batches.get(name=batch_id) dest = getattr(job, "dest", None) - inlined = getattr(dest, "inlined_responses", None) if dest is not None else None - if inlined: - for item in inlined: - yield _parse_gemini_inline_response(item) - return file_name = getattr(dest, "file_name", None) if dest is not None else None - if file_name: - content = self.genai_client.files.download(file=file_name) - text = ( - content.decode("utf-8") - if isinstance(content, (bytes, bytearray)) - else str(content) - ) - for line in text.splitlines(): - stripped = line.strip() - if stripped: - yield _parse_gemini_result_line(stripped) - - -def _extract_gemini_text(response: Any) -> str: - """Extract the assistant text from a Gemini response (dict from file, or SDK object).""" - if isinstance(response, dict): - for candidate in response.get("candidates") or []: - content = (candidate or {}).get("content") or {} - chunks = [ - str(part.get("text")) - for part in content.get("parts") or [] - if isinstance(part, dict) and part.get("text") - ] - if chunks: - return "".join(chunks) - return "" - text = getattr(response, "text", None) - return str(text) if text else "" - - -def _parse_gemini_inline_response(item: Any) -> BatchResult: - """Parse one ``dest.inlined_responses`` entry into a ``BatchResult``.""" - key = str(getattr(item, "key", "") or "") - error = getattr(item, "error", None) - response = getattr(item, "response", None) - if error is not None or response is None: - return BatchResult( - key, - BatchItemStatus.ERRORED, - error=str(error) if error is not None else "no response", + if not file_name: + return + content = self.genai_client.files.download(file=file_name) + text = ( + content.decode("utf-8") + if isinstance(content, (bytes, bytearray)) + else str(content) ) - return BatchResult(key, BatchItemStatus.SUCCEEDED, text=_extract_gemini_text(response)) + for line in text.splitlines(): + stripped = line.strip() + if stripped: + yield _parse_gemini_result_line(stripped) + + +def _extract_gemini_text(response: dict[str, Any]) -> str: + """Extract the assistant text from a batch ``response`` dict (parsed from JSONL).""" + for candidate in response.get("candidates") or []: + content = (candidate or {}).get("content") or {} + chunks = [ + str(part.get("text")) + for part in content.get("parts") or [] + if isinstance(part, dict) and part.get("text") + ] + if chunks: + return "".join(chunks) + return "" def _parse_gemini_result_line(line: str) -> BatchResult: @@ -343,7 +341,9 @@ def _parse_gemini_result_line(line: str) -> BatchResult: if error is not None or response is None: message = json.dumps(error) if error is not None else "no response" return BatchResult(key, BatchItemStatus.ERRORED, error=message) - return BatchResult(key, BatchItemStatus.SUCCEEDED, text=_extract_gemini_text(response)) + return BatchResult( + key, BatchItemStatus.SUCCEEDED, text=_extract_gemini_text(response) + ) def _extract_gemini_error_details(response: object) -> str | None: diff --git a/tests/prkit/core/model_clients/test_batch.py b/tests/prkit/core/model_clients/test_batch.py index e81251a..a178a5b 100644 --- a/tests/prkit/core/model_clients/test_batch.py +++ b/tests/prkit/core/model_clients/test_batch.py @@ -54,13 +54,23 @@ def test_terminal_states(self): def test_status_is_terminal(self): assert BatchStatus("b", BatchState.COMPLETED, "openai", "completed").is_terminal - assert not BatchStatus("b", BatchState.PENDING, "openai", "validating").is_terminal + assert not BatchStatus( + "b", BatchState.PENDING, "openai", "validating" + ).is_terminal class TestBaseDefaults: def test_batch_methods_raise_not_implemented(self): class Dummy(BaseModelClient): - def response(self, input, image_paths=None, response_format=None, *, instructions=None, **kwargs): + def response( + self, + input, + image_paths=None, + response_format=None, + *, + instructions=None, + **kwargs, + ): return "" d = Dummy("dummy-model") @@ -78,7 +88,11 @@ class TestFreeTextBuilders: def test_openai_shape(self): client = _openai_client() req = client.build_batch_request( - request_id="req_1", input="hello", instructions="", max_output_tokens=100, temperature=0.0 + request_id="req_1", + input="hello", + instructions="", + max_output_tokens=100, + temperature=0.0, ) assert req["custom_id"] == "req_1" assert req["method"] == "POST" @@ -94,7 +108,11 @@ def test_openai_shape(self): def test_openai_o_family_drops_temperature_sets_reasoning(self): client = _openai_client("o3") body = client.build_batch_request( - request_id="r", input="x", instructions="", max_output_tokens=50, temperature=0.7 + request_id="r", + input="x", + instructions="", + max_output_tokens=50, + temperature=0.7, )["body"] assert "temperature" not in body assert body["reasoning"] == {"effort": "medium"} @@ -102,7 +120,11 @@ def test_openai_o_family_drops_temperature_sets_reasoning(self): def test_anthropic_shape(self): client = _anthropic_client() params = client.build_batch_request( - request_id="req_2", input="hi", instructions="", max_output_tokens=64, temperature=0.2 + request_id="req_2", + input="hi", + instructions="", + max_output_tokens=64, + temperature=0.2, )["params"] assert params["model"] == "claude-opus-4-8" assert params["max_tokens"] == 64 @@ -113,13 +135,19 @@ def test_anthropic_shape(self): def test_anthropic_defaults_max_tokens_when_none(self): client = _anthropic_client() - params = client.build_batch_request(request_id="r", input="x", instructions="")["params"] + params = client.build_batch_request(request_id="r", input="x", instructions="")[ + "params" + ] assert params["max_tokens"] == 1024 def test_gemini_shape(self): client = _gemini_client() req = client.build_batch_request( - request_id="req_3", input="q", instructions="", max_output_tokens=32, temperature=0.5 + request_id="req_3", + input="q", + instructions="", + max_output_tokens=32, + temperature=0.5, ) assert req["key"] == "req_3" request = req["request"] @@ -203,7 +231,11 @@ def test_retrieve_extracts_output_text(self): def test_retrieve_flags_http_error(self): client = _openai_client() line = json.dumps( - {"custom_id": "r2", "response": {"status_code": 500, "body": {}}, "error": None} + { + "custom_id": "r2", + "response": {"status_code": 500, "body": {}}, + "error": None, + } ) client.client.batches.retrieve.return_value = MagicMock( output_file_id="out_1", error_file_id=None @@ -259,14 +291,32 @@ def test_retrieve_errored(self): class TestGeminiLifecycle: - def test_submit(self): + def test_submit_uploads_keyed_jsonl_file(self): client = _gemini_client() + uploaded = MagicMock() + uploaded.name = "files/in" + client.genai_client.files.upload.return_value = uploaded job = MagicMock() job.name = "batches/xyz" client.genai_client.batches.create.return_value = job - assert client.submit_batch([{"key": "r1", "request": {}}]) == "batches/xyz" - _, kwargs = client.genai_client.batches.create.call_args - assert kwargs["model"] == "gemini-3.5-flash" + + result = client.submit_batch( + [ + {"key": "r1", "request": {"contents": []}}, + {"key": "r2", "request": {"contents": []}}, + ] + ) + + assert result == "batches/xyz" + # The batch must be created from the uploaded file, not an inline list. + _, create_kwargs = client.genai_client.batches.create.call_args + assert create_kwargs["model"] == "gemini-3.5-flash" + assert create_kwargs["src"] == "files/in" + # The uploaded payload must be keyed JSONL so results correlate by key. + _, upload_kwargs = client.genai_client.files.upload.call_args + payload = upload_kwargs["file"].getvalue().decode("utf-8") + lines = [json.loads(line) for line in payload.splitlines()] + assert [line["key"] for line in lines] == ["r1", "r2"] def test_poll(self): client = _gemini_client() @@ -279,31 +329,18 @@ def test_poll(self): assert status.state == BatchState.COMPLETED assert status.output_ref == "files/out" - def test_retrieve_inline(self): - client = _gemini_client() - item = MagicMock() - item.key = "r1" - item.error = None - item.response = MagicMock(text="GENTEXT") - job = MagicMock() - job.dest = MagicMock() - job.dest.inlined_responses = [item] - client.genai_client.batches.get.return_value = job - results = list(client.retrieve_batch_results("batches/xyz")) - assert results[0].custom_id == "r1" - assert results[0].text == "GENTEXT" - def test_retrieve_file(self): client = _gemini_client() job = MagicMock() job.dest = MagicMock() - job.dest.inlined_responses = None job.dest.file_name = "files/out" client.genai_client.batches.get.return_value = job line = json.dumps( { "key": "r2", - "response": {"candidates": [{"content": {"parts": [{"text": "FILETEXT"}]}}]}, + "response": { + "candidates": [{"content": {"parts": [{"text": "FILETEXT"}]}}] + }, } ) client.genai_client.files.download.return_value = line.encode("utf-8") @@ -311,6 +348,16 @@ def test_retrieve_file(self): assert results[0].custom_id == "r2" assert results[0].text == "FILETEXT" + def test_retrieve_without_output_file_yields_nothing(self): + # No destination file => no results, rather than silently miscorrelated ones. + client = _gemini_client() + job = MagicMock() + job.dest = MagicMock() + job.dest.file_name = None + client.genai_client.batches.get.return_value = job + assert list(client.retrieve_batch_results("batches/xyz")) == [] + client.genai_client.files.download.assert_not_called() + class TestSyncBatchParity: """A batch request body must equal the body sync ``response()`` would build.""" @@ -327,7 +374,11 @@ def test_openai_body_parity(self): extra={"temperature": 0.0}, ) batch_body = client.build_batch_request( - request_id="r", input="P", instructions="", max_output_tokens=123, temperature=0.0 + request_id="r", + input="P", + instructions="", + max_output_tokens=123, + temperature=0.0, )["body"] assert batch_body == sync_body @@ -343,6 +394,10 @@ def test_anthropic_params_parity(self): extra={"temperature": 0.2}, ) batch_params = client.build_batch_request( - request_id="r", input="P", instructions="", max_output_tokens=64, temperature=0.2 + request_id="r", + input="P", + instructions="", + max_output_tokens=64, + temperature=0.2, )["params"] assert batch_params == sync_params From b7e70527bcc9a058b10dc6b031bd3a633794ac0f Mon Sep 17 00:00:00 2001 From: Yinghuan Zhang Date: Mon, 15 Jun 2026 12:24:07 -0400 Subject: [PATCH 5/8] Add gated live smoke test for Gemini batch correlation Integration test (skipped unless GEMINI_API_KEY/GOOGLE_API_KEY set) that submits a real batch with distinct ids + unique answers and asserts each answer lands under the correct custom_id -- the identity check offline mocks can't make. Deterministic arithmetic prompts + thinking-aware token budget so it validates correlation, not model instruction-following. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../core/model_clients/test_batch_live.py | 100 ++++++++++++++++++ 1 file changed, 100 insertions(+) create mode 100644 tests/prkit/core/model_clients/test_batch_live.py diff --git a/tests/prkit/core/model_clients/test_batch_live.py b/tests/prkit/core/model_clients/test_batch_live.py new file mode 100644 index 0000000..50885a4 --- /dev/null +++ b/tests/prkit/core/model_clients/test_batch_live.py @@ -0,0 +1,100 @@ +"""Live, gated smoke test for the Gemini file-based batch path. + +This is the check the offline suite *cannot* make: it submits a real batch and +confirms each answer lands under the **correct** ``custom_id`` (identity), not +merely that some text came back. It is the difference between a working +file-based correlation and a silent positional misalignment. + +Skipped unless a Gemini API key is present, so default ``pytest`` runs ignore +it. Run explicitly with:: + + GEMINI_API_KEY=... .venv/bin/pytest \ + tests/prkit/core/model_clients/test_batch_live.py -m integration -v -s + +Polling timeout is configurable via ``PRKIT_BATCH_TIMEOUT_SECONDS`` (default +1800s); batch jobs are asynchronous and may take minutes to complete. +""" + +from __future__ import annotations + +import os +import time + +import pytest + +from prkit.core.model_clients.batch_types import BatchItemStatus + +pytestmark = [pytest.mark.integration, pytest.mark.slow] + +_HAS_KEY = bool(os.environ.get("GEMINI_API_KEY") or os.environ.get("GOOGLE_API_KEY")) + +# request_id -> (prompt, unique expected answer). Distinct arithmetic gives +# deterministic, non-overlapping answers that any model returns reliably, so the +# test validates *correlation* (the right answer under the right id) rather than +# the model's instruction-following. Answers are chosen so none is a substring of +# another, which is what makes the "leak" check below meaningful. +_CASES = { + "q-alpha": ("What is 40 + 2? Reply with only the number.", "42"), + "q-bravo": ("What is 50 + 3? Reply with only the number.", "53"), + "q-charlie": ("What is 60 + 4? Reply with only the number.", "64"), +} + + +@pytest.mark.skipif(not _HAS_KEY, reason="No GEMINI_API_KEY/GOOGLE_API_KEY set") +def test_gemini_batch_correlates_each_answer_to_its_request(): + from prkit.core.model_clients.gemini import GeminiModel + + client = GeminiModel(os.environ.get("PRKIT_BATCH_MODEL", "gemini-3.5-flash")) + + requests = [ + client.build_batch_request( + request_id=request_id, + input=prompt, + instructions="", + # gemini-3.5-flash is a thinking model: the budget must cover hidden + # reasoning tokens (observed ~100-500) plus the short answer. + max_output_tokens=2048, + temperature=0.0, + ) + for request_id, (prompt, _expected) in _CASES.items() + ] + + batch_id = client.submit_batch( + requests, metadata={"display_name": "prkit-live-smoke"} + ) + assert batch_id + + timeout = float(os.environ.get("PRKIT_BATCH_TIMEOUT_SECONDS", "1800")) + deadline = time.monotonic() + timeout + status = client.poll_batch(batch_id) + while not status.is_terminal: + if time.monotonic() > deadline: + pytest.fail( + f"Batch {batch_id} not terminal after {timeout}s (last={status.raw_status})" + ) + time.sleep(20) + status = client.poll_batch(batch_id) + + results = {r.custom_id: r for r in client.retrieve_batch_results(batch_id)} + + # Every request must come back, keyed by exactly the id we submitted. + assert set(results) == set( + _CASES + ), f"custom_id mismatch: {set(results)} != {set(_CASES)}" + + for request_id, (_prompt, expected) in _CASES.items(): + result = results[request_id] + assert ( + result.status is BatchItemStatus.SUCCEEDED + ), f"{request_id}: {result.error}" + text = (result.text or "").upper() + # The answer for this id must contain ITS word and none of the others' — + # this is what catches a positional misalignment. + assert ( + expected in text + ), f"{request_id}: expected {expected!r} in {result.text!r}" + for other_id, (_p, other_word) in _CASES.items(): + if other_id != request_id: + assert ( + other_word not in text + ), f"{request_id} answer leaked {other_word!r}: {result.text!r}" From b6504a53ba463d223f49e47f9fcfd8a38d5e850f Mon Sep 17 00:00:00 2001 From: Yinghuan Zhang Date: Mon, 15 Jun 2026 13:03:41 -0400 Subject: [PATCH 6/8] Apply black formatting to anthropic/openai clients and uq conftest Co-Authored-By: Claude Opus 4.8 (1M context) --- src/prkit/core/model_clients/anthropic.py | 8 ++++++-- src/prkit/core/model_clients/openai.py | 10 ++++++++-- tests/uq/conftest.py | 5 ++++- 3 files changed, 18 insertions(+), 5 deletions(-) diff --git a/src/prkit/core/model_clients/anthropic.py b/src/prkit/core/model_clients/anthropic.py index 4cc7cce..700c325 100644 --- a/src/prkit/core/model_clients/anthropic.py +++ b/src/prkit/core/model_clients/anthropic.py @@ -359,7 +359,9 @@ def _build_batch_request( input=input, instructions=instructions, image_paths=image_paths, - max_output_tokens=max_output_tokens if max_output_tokens is not None else 1024, + max_output_tokens=( + max_output_tokens if max_output_tokens is not None else 1024 + ), response_format=None, extra=extra, ) @@ -457,7 +459,9 @@ def _build_batch_structured_request( input=user_prompt, instructions=None, image_paths=image_paths, - max_output_tokens=max_output_tokens if max_output_tokens is not None else 4096, + max_output_tokens=( + max_output_tokens if max_output_tokens is not None else 4096 + ), response_format=plan.response_format or {}, ) return {"custom_id": request_id, "params": params} diff --git a/src/prkit/core/model_clients/openai.py b/src/prkit/core/model_clients/openai.py index fe35986..6c96c74 100644 --- a/src/prkit/core/model_clients/openai.py +++ b/src/prkit/core/model_clients/openai.py @@ -570,8 +570,14 @@ def _parse_openai_result_line(line: str, *, force_error: bool = False) -> BatchR error = obj.get("error") response = obj.get("response") if isinstance(obj.get("response"), dict) else None status_code = response.get("status_code") if response else None - if force_error or error is not None or (status_code is not None and status_code != 200): - message = json.dumps(error) if error is not None else f"status_code={status_code}" + if ( + force_error + or error is not None + or (status_code is not None and status_code != 200) + ): + message = ( + json.dumps(error) if error is not None else f"status_code={status_code}" + ) return BatchResult(custom_id, BatchItemStatus.ERRORED, error=message) body = response.get("body") if response else None return BatchResult( diff --git a/tests/uq/conftest.py b/tests/uq/conftest.py index 77891ea..28a6e80 100644 --- a/tests/uq/conftest.py +++ b/tests/uq/conftest.py @@ -1,6 +1,9 @@ import importlib.util -_uq_available = importlib.util.find_spec("uncertainty_quantification_physical_reasoning") is not None +_uq_available = ( + importlib.util.find_spec("uncertainty_quantification_physical_reasoning") + is not None +) collect_ignore: list[str] = [] if not _uq_available: From e0b9c470aa80ded651755babaae007aea20f1930 Mon Sep 17 00:00:00 2001 From: Yinghuan Zhang Date: Mon, 15 Jun 2026 14:27:45 -0400 Subject: [PATCH 7/8] Guard Gemini batch upload file name against None for mypy files.upload(...).name is typed str | None, but batches.create(src=...) requires a non-None value. Raise a clear RuntimeError if the upload returns no name instead of passing Optional[str] through. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/prkit/core/model_clients/gemini.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/src/prkit/core/model_clients/gemini.py b/src/prkit/core/model_clients/gemini.py index d7f8b13..67cded0 100644 --- a/src/prkit/core/model_clients/gemini.py +++ b/src/prkit/core/model_clients/gemini.py @@ -273,9 +273,14 @@ def submit_batch( file=io.BytesIO(payload.encode("utf-8")), config=types.UploadFileConfig(mime_type="jsonl", display_name=display_name), ) + uploaded_name = uploaded.name + if uploaded_name is None: + raise RuntimeError( + "Gemini File API upload returned no file name; cannot submit batch" + ) job = self.genai_client.batches.create( model=self.model, - src=uploaded.name, + src=uploaded_name, config={"display_name": display_name}, ) return str(job.name) From 15378fa090c38a9e74c48ee6fa16b4a874917088 Mon Sep 17 00:00:00 2001 From: Yinghuan Zhang Date: Mon, 15 Jun 2026 14:38:53 -0400 Subject: [PATCH 8/8] Add mypy and pytest pre-commit hooks mirroring CI The pre-commit config ran ruff/black/whitespace but omitted mypy and pytest, so type and test failures only surfaced in CI -- as the Gemini batch Optional[str] mypy error did on this branch. Add local hooks that invoke the project venv to run `mypy src/prkit` and `pytest tests/prkit`, matching .github/workflows/ci.yml, so both fail at commit time instead. Scoped via files: filters so they fire only when src/prkit or tests/prkit is staged. Co-Authored-By: Claude Opus 4.8 (1M context) --- .pre-commit-config.yaml | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 59466ee..9ed3134 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -18,3 +18,23 @@ repos: hooks: - id: end-of-file-fixer - id: trailing-whitespace + + # Local hooks mirror the CI checks (.github/workflows/ci.yml) so type and + # test failures surface at commit time instead of in CI. They invoke the + # project venv's interpreter (.venv/bin/python) directly because mypy and + # pytest need the installed dependencies, which an isolated pre-commit env + # would not have, and so they work whether or not the venv is activated. + - repo: local + hooks: + - id: mypy + name: mypy (src/prkit) + entry: .venv/bin/python -m mypy src/prkit + language: system + files: ^src/prkit/ + pass_filenames: false + - id: pytest + name: pytest (tests/prkit) + entry: .venv/bin/python -m pytest tests/prkit -q + language: system + files: ^(src/prkit|tests/prkit)/ + pass_filenames: false