This is an automated email from the ASF dual-hosted git repository.
wenjin272 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/flink-agents.git
The following commit(s) were added to refs/heads/main by this push:
new 88fab2b39 [api][python][integrations] Share the native output model
helper across chat model connections (#1277)
88fab2b39 is described below
commit 88fab2b395da7d7150c1c4969ccbf0d07c1995ab
Author: Weiqing Yang <[email protected]>
AuthorDate: Mon Oct 19 19:32:42 2026 -0700
[api][python][integrations] Share the native output model helper across
chat model connections (#1277)
Generated-by: GitHub Copilot CLI 1.0.98 (Claude Opus 5.5)
---
python/flink_agents/api/agents/types.py | 25 ++++++++++++++++
.../chat_models/anthropic/anthropic_chat_model.py | 27 +++---------------
.../chat_models/azure/azure_openai_chat_model.py | 28 ++++--------------
.../chat_models/dashscope_chat_model.py | 31 ++++++--------------
.../integrations/chat_models/ollama_chat_model.py | 32 ++++++---------------
.../chat_models/openai/openai_chat_model.py | 27 +++---------------
.../chat_models/vllm/vllm_chat_model.py | 5 ++--
.../chat_models/watsonx/watsonx_chat_model.py | 33 ++++++----------------
8 files changed, 66 insertions(+), 142 deletions(-)
diff --git a/python/flink_agents/api/agents/types.py
b/python/flink_agents/api/agents/types.py
index cc1528356..1f8e9a29c 100644
--- a/python/flink_agents/api/agents/types.py
+++ b/python/flink_agents/api/agents/types.py
@@ -68,6 +68,31 @@ class OutputSchema(BaseModel):
return self
+def native_output_model(output_schema: OutputSchema | None) -> type[BaseModel]
| None:
+ """The Pydantic model an output schema would be sent natively as, or
``None``.
+
+ Kept apart from ``render_output_schema`` because a connection makes two
checks
+ before rendering: whether native structured output can apply, and whether
the
+ schema conflicts with a caller-supplied response format. Neither check
renders
+ the schema. Rendering raises on a schema it cannot express, so the check
would
+ report a render failure instead of its own answer.
+
+ Args:
+ output_schema: The output schema requested for the response, if any.
+
+ Returns:
+ The ``BaseModel`` subclass the schema holds. ``None`` covers both no
schema at
+ all and a ``RowTypeInfo``, which has no native translation and keeps
the
+ prompt-engineering fallback.
+ """
+ model = (
+ output_schema.output_schema if isinstance(output_schema, OutputSchema)
else None
+ )
+ if not (isinstance(model, type) and issubclass(model, BaseModel)):
+ return None
+ return model
+
+
def render_output_schema(
model: type[BaseModel], render: Callable[[type[BaseModel]], dict[str, Any]]
) -> dict[str, Any]:
diff --git
a/python/flink_agents/integrations/chat_models/anthropic/anthropic_chat_model.py
b/python/flink_agents/integrations/chat_models/anthropic/anthropic_chat_model.py
index 7264d5a3c..417dfc414 100644
---
a/python/flink_agents/integrations/chat_models/anthropic/anthropic_chat_model.py
+++
b/python/flink_agents/integrations/chat_models/anthropic/anthropic_chat_model.py
@@ -24,12 +24,13 @@ from typing import Any, Dict, List, Mapping, Sequence
from anthropic import Anthropic, transform_schema
from anthropic._types import NOT_GIVEN
from anthropic.types import MessageParam, TextBlockParam, ToolParam
-from pydantic import BaseModel, Field, PrivateAttr
+from pydantic import Field, PrivateAttr
from typing_extensions import override
from flink_agents.api.agents.types import (
OutputSchema,
contains_open_object,
+ native_output_model,
render_output_schema,
)
from flink_agents.api.chat_message import (
@@ -371,26 +372,6 @@ def _supports_sampling_params(effective_model: str | None)
-> bool:
return effective_model not in _SAMPLING_UNSUPPORTED_MODELS
-def _native_output_model(output_schema: Any) -> type[BaseModel] | None:
- """The model a schema translates natively to, or ``None`` where none
applies.
-
- ``None`` covers both no schema at all and a ``RowTypeInfo``, which has no
native
- translation and keeps the prompt-engineering fallback.
-
- Separate from the render below because the feasibility check has to know
whether a
- schema would be sent without rendering it, and rendering raises on a
schema it
- cannot express.
- """
- if output_schema is None:
- return None
- model = (
- output_schema.output_schema if isinstance(output_schema, OutputSchema)
else None
- )
- if not (isinstance(model, type) and issubclass(model, BaseModel)):
- return None
- return model
-
-
def _native_output_config(output_schema: Any) -> Dict[str, Any] | None:
"""Build the Anthropic ``output_config`` for a native structured-output
request.
@@ -408,7 +389,7 @@ def _native_output_config(output_schema: Any) -> Dict[str,
Any] | None:
but declares no fields is sent as it is, leaving the provider to accept or
refuse
the document it receives.
"""
- model = _native_output_model(output_schema)
+ model = native_output_model(output_schema)
if model is None:
return None
return {
@@ -560,7 +541,7 @@ class AnthropicChatModelConnection(BaseChatModelConnection):
admits no undeclared keys, and the caller supplied no
``output_config`` of
its own.
"""
- model = _native_output_model(output_schema)
+ model = native_output_model(output_schema)
return (
model is not None
and (model_kwargs is None or "output_config" not in model_kwargs)
diff --git
a/python/flink_agents/integrations/chat_models/azure/azure_openai_chat_model.py
b/python/flink_agents/integrations/chat_models/azure/azure_openai_chat_model.py
index 7820f8135..e9776cbe6 100644
---
a/python/flink_agents/integrations/chat_models/azure/azure_openai_chat_model.py
+++
b/python/flink_agents/integrations/chat_models/azure/azure_openai_chat_model.py
@@ -26,12 +26,13 @@ from openai import NOT_GIVEN, AzureOpenAI
# has existed at this path since the structured-output support in openai
1.66.3 (the
# pinned minimum). A future openai bump that moves it will fail loudly on
import here.
from openai.lib._pydantic import to_strict_json_schema
-from pydantic import BaseModel, Field, PrivateAttr
+from pydantic import Field, PrivateAttr
from typing_extensions import override
from flink_agents.api.agents.types import (
OutputSchema,
contains_open_object,
+ native_output_model,
render_output_schema,
)
from flink_agents.api.chat_message import ChatMessage
@@ -99,25 +100,6 @@ _MIN_STRUCTURED_OUTPUT_API_VERSION = "2024-08-01"
_API_VERSION_DATE_PREFIX = re.compile(r"^\d{4}-\d{2}-\d{2}", re.ASCII)
-def _native_output_model(output_schema: Any) -> type[BaseModel] | None:
- """The model a schema translates natively to, or ``None`` where none
applies.
-
- ``None`` covers both no schema at all and a ``RowTypeInfo``, which has no
native
- translation and keeps the prompt-engineering fallback.
-
- Separate from the render below because the caller-conflict check needs to
know
- whether a schema will be sent, and under what name, before anything is
rendered.
- """
- if output_schema is None:
- return None
- model = (
- output_schema.output_schema if isinstance(output_schema, OutputSchema)
else None
- )
- if not (isinstance(model, type) and issubclass(model, BaseModel)):
- return None
- return model
-
-
def _native_response_format(output_schema: Any) -> Dict[str, Any] | None:
"""Build the ``response_format`` for a native structured-output request.
@@ -131,7 +113,7 @@ def _native_response_format(output_schema: Any) ->
Dict[str, Any] | None:
but declares no fields is sent as it is, leaving the provider to accept or
refuse
the document it receives.
"""
- model = _native_output_model(output_schema)
+ model = native_output_model(output_schema)
if model is None:
return None
return {
@@ -344,7 +326,7 @@ class
AzureOpenAIChatModelConnection(BaseChatModelConnection):
admits no undeclared keys, and the configured api-version reaches
the
structured-output floor.
"""
- model = _native_output_model(output_schema)
+ model = native_output_model(output_schema)
return (
model is not None
and self._api_version_supports_structured_output()
@@ -421,7 +403,7 @@ class
AzureOpenAIChatModelConnection(BaseChatModelConnection):
# INFEASIBLE. Model capability is not asked: a caller that hands this
# connection a schema has already chosen to send it.
if self._can_apply_native_structured_output(output_schema, tools,
raw_kwargs):
- native_model = _native_output_model(output_schema)
+ native_model = native_output_model(output_schema)
# Tested before the schema is rendered. A caller who supplies both
a schema
# and a response_format has a conflict to resolve whatever the
schema turns
# out to render to, and reporting a render failure instead would
describe
diff --git
a/python/flink_agents/integrations/chat_models/dashscope_chat_model.py
b/python/flink_agents/integrations/chat_models/dashscope_chat_model.py
index d5404b309..d5b404ac6 100644
--- a/python/flink_agents/integrations/chat_models/dashscope_chat_model.py
+++ b/python/flink_agents/integrations/chat_models/dashscope_chat_model.py
@@ -21,10 +21,14 @@ import uuid
from typing import Any, Dict, List, Mapping, Sequence
from dashscope import Generation
-from pydantic import BaseModel, Field
+from pydantic import Field
from typing_extensions import override
-from flink_agents.api.agents.types import OutputSchema, render_output_schema
+from flink_agents.api.agents.types import (
+ OutputSchema,
+ native_output_model,
+ render_output_schema,
+)
from flink_agents.api.chat_message import (
ChatMessage,
MessageRole,
@@ -68,23 +72,6 @@ _NATIVE_STRUCTURED_OUTPUT_MODELS = frozenset(
)
-def _native_output_model(
- output_schema: OutputSchema | None,
-) -> type[BaseModel] | None:
- """The model a schema translates natively to, or ``None`` where none
applies.
-
- ``None`` covers both no schema at all and a ``RowTypeInfo``, which has no
native
- translation and keeps the prompt-engineering fallback.
-
- Separate from the render below because the caller-conflict check needs to
know
- whether a schema will be sent, and under what name, before anything is
rendered.
- """
- schema = getattr(output_schema, "output_schema", None)
- if not (isinstance(schema, type) and issubclass(schema, BaseModel)):
- return None
- return schema
-
-
def _native_response_format(
output_schema: OutputSchema | None,
) -> Dict[str, Any] | None:
@@ -100,7 +87,7 @@ def _native_response_format(
but declares no fields is sent as it is, leaving the provider to accept or
refuse
the document it receives.
"""
- model = _native_output_model(output_schema)
+ model = native_output_model(output_schema)
if model is None:
return None
return {
@@ -273,7 +260,7 @@ class DashScopeChatModelConnection(BaseChatModelConnection):
bool
``True`` if ``output_schema`` wraps a ``BaseModel`` subclass.
"""
- return _native_output_model(output_schema) is not None
+ return native_output_model(output_schema) is not None
def chat(
self,
@@ -339,7 +326,7 @@ class DashScopeChatModelConnection(BaseChatModelConnection):
# schema turns out to render to, and reporting a render failure
instead
# would describe the wrong problem. The name is read off the model
class,
# so this needs no rendered document.
- native_model = _native_output_model(output_schema)
+ native_model = native_output_model(output_schema)
if "response_format" in kwargs:
msg = (
f"The {native_model.__name__} output schema is sent as "
diff --git a/python/flink_agents/integrations/chat_models/ollama_chat_model.py
b/python/flink_agents/integrations/chat_models/ollama_chat_model.py
index b350e7c48..3b157f7a4 100644
--- a/python/flink_agents/integrations/chat_models/ollama_chat_model.py
+++ b/python/flink_agents/integrations/chat_models/ollama_chat_model.py
@@ -21,10 +21,14 @@ import uuid
from typing import Any, Dict, List, Literal, Mapping, Sequence
from ollama import Client, Image, Message
-from pydantic import BaseModel, Field
+from pydantic import Field
from typing_extensions import override
-from flink_agents.api.agents.types import OutputSchema, render_output_schema
+from flink_agents.api.agents.types import (
+ OutputSchema,
+ native_output_model,
+ render_output_schema,
+)
from flink_agents.api.chat_message import (
Base64Source,
ChatMessage,
@@ -49,26 +53,6 @@ DEFAULT_CONTEXT_WINDOW = 2048
DEFAULT_REQUEST_TIMEOUT = 30.0
-def _native_output_model(output_schema: Any) -> type[BaseModel] | None:
- """The model a schema translates natively to, or ``None`` where none
applies.
-
- ``None`` covers both no schema at all and a ``RowTypeInfo``, which has no
native
- translation and keeps the prompt-engineering fallback.
-
- Separate from the render below because the feasibility check has to know
whether a
- schema would be sent without rendering it, and rendering raises on a
schema it
- cannot express.
- """
- if output_schema is None:
- return None
- model = (
- output_schema.output_schema if isinstance(output_schema, OutputSchema)
else None
- )
- if not (isinstance(model, type) and issubclass(model, BaseModel)):
- return None
- return model
-
-
def _native_format(output_schema: Any) -> Dict[str, Any] | None:
"""Build the Ollama ``format`` payload for a native structured-output
request.
@@ -77,7 +61,7 @@ def _native_format(output_schema: Any) -> Dict[str, Any] |
None:
prompt-engineering fallback. A schema that cannot be rendered raises
through
``render_output_schema``, naming the model, like the other connections.
"""
- model = _native_output_model(output_schema)
+ model = native_output_model(output_schema)
if model is None:
return None
return render_output_schema(model, lambda m: m.model_json_schema())
@@ -199,7 +183,7 @@ class OllamaChatModelConnection(BaseChatModelConnection):
bool
``True`` if ``output_schema`` wraps a ``BaseModel`` subclass.
"""
- return _native_output_model(output_schema) is not None
+ return native_output_model(output_schema) is not None
def chat(
self,
diff --git
a/python/flink_agents/integrations/chat_models/openai/openai_chat_model.py
b/python/flink_agents/integrations/chat_models/openai/openai_chat_model.py
index e7df77466..aa7214cb3 100644
--- a/python/flink_agents/integrations/chat_models/openai/openai_chat_model.py
+++ b/python/flink_agents/integrations/chat_models/openai/openai_chat_model.py
@@ -26,12 +26,13 @@ from openai import NOT_GIVEN, OpenAI
# has existed at this path since the structured-output support in openai
1.66.3.
# A future openai bump that moves it will fail loudly on import here.
from openai.lib._pydantic import to_strict_json_schema
-from pydantic import BaseModel, Field, PrivateAttr
+from pydantic import Field, PrivateAttr
from typing_extensions import override
from flink_agents.api.agents.types import (
OutputSchema,
contains_open_object,
+ native_output_model,
render_output_schema,
)
from flink_agents.api.chat_message import ChatMessage
@@ -89,26 +90,6 @@ _NATIVE_STRUCTURED_OUTPUT_MODELS = frozenset(
)
-def _native_output_model(output_schema: Any) -> type[BaseModel] | None:
- """The model a schema translates natively to, or ``None`` where none
applies.
-
- ``None`` covers both no schema at all and a ``RowTypeInfo``, which has no
native
- translation and keeps the prompt-engineering fallback.
-
- Separate from the render below because the feasibility check has to know
whether a
- schema would be sent without rendering it, and rendering raises on a
schema it
- cannot express.
- """
- if output_schema is None:
- return None
- model = (
- output_schema.output_schema if isinstance(output_schema, OutputSchema)
else None
- )
- if not (isinstance(model, type) and issubclass(model, BaseModel)):
- return None
- return model
-
-
def _native_response_format(output_schema: Any) -> Dict[str, Any] | None:
"""Build the OpenAI ``response_format`` for a native structured-output
request.
@@ -122,7 +103,7 @@ def _native_response_format(output_schema: Any) ->
Dict[str, Any] | None:
but declares no fields is sent as it is, leaving the provider to accept or
refuse
the document it receives.
"""
- model = _native_output_model(output_schema)
+ model = native_output_model(output_schema)
if model is None:
return None
return {
@@ -317,7 +298,7 @@ class OpenAIChatModelConnection(BaseChatModelConnection):
``True`` if ``output_schema`` wraps a ``BaseModel`` subclass whose
schema
admits no undeclared keys.
"""
- model = _native_output_model(output_schema)
+ model = native_output_model(output_schema)
return model is not None and not contains_open_object(model)
def chat(
diff --git
a/python/flink_agents/integrations/chat_models/vllm/vllm_chat_model.py
b/python/flink_agents/integrations/chat_models/vllm/vllm_chat_model.py
index a427801d1..db5bc464e 100644
--- a/python/flink_agents/integrations/chat_models/vllm/vllm_chat_model.py
+++ b/python/flink_agents/integrations/chat_models/vllm/vllm_chat_model.py
@@ -19,12 +19,11 @@ from typing import Any, List, Mapping
from typing_extensions import override
-from flink_agents.api.agents.types import OutputSchema
+from flink_agents.api.agents.types import OutputSchema, native_output_model
from flink_agents.api.tools.tool import Tool
from flink_agents.integrations.chat_models.openai.openai_chat_model import (
OpenAIChatModelConnection,
OpenAIChatModelSetup,
- _native_output_model,
)
DEFAULT_VLLM_API_BASE_URL = "http://localhost:8000/v1";
@@ -104,7 +103,7 @@ class VLLMChatModelConnection(OpenAIChatModelConnection):
rules, and its guided decoding honours ``additionalProperties``, so a
map
keeps its value type. Only a ``BaseModel`` subclass has a native
translation.
"""
- return _native_output_model(output_schema) is not None
+ return native_output_model(output_schema) is not None
class VLLMChatModelSetup(OpenAIChatModelSetup):
diff --git
a/python/flink_agents/integrations/chat_models/watsonx/watsonx_chat_model.py
b/python/flink_agents/integrations/chat_models/watsonx/watsonx_chat_model.py
index f4cbfb379..1a97cebd7 100644
--- a/python/flink_agents/integrations/chat_models/watsonx/watsonx_chat_model.py
+++ b/python/flink_agents/integrations/chat_models/watsonx/watsonx_chat_model.py
@@ -28,10 +28,14 @@ import httpx
from ibm_watsonx_ai import APIClient, Credentials
from ibm_watsonx_ai.foundation_models import ModelInference
from ibm_watsonx_ai.wml_client_error import ApiRequestFailure
-from pydantic import BaseModel, Field, PrivateAttr
+from pydantic import Field, PrivateAttr
from typing_extensions import override
-from flink_agents.api.agents.types import OutputSchema, render_output_schema
+from flink_agents.api.agents.types import (
+ OutputSchema,
+ native_output_model,
+ render_output_schema,
+)
from flink_agents.api.chat_message import (
ChatMessage,
MessageRole,
@@ -167,25 +171,6 @@ def _convert_to_watsonx_tool_call(tool_call:
ToolCallBlock) -> Dict[str, Any]:
}
-def _native_output_model(
- output_schema: OutputSchema | None,
-) -> type[BaseModel] | None:
- """The model a schema translates natively to, or ``None`` where none
applies.
-
- ``None`` covers both no schema at all and a ``RowTypeInfo``, which has no
native
- translation here and keeps the prompt-engineering fallback.
-
- Separate from the render below because the caller-conflict check needs to
know
- whether a schema will be sent, and under what name, before anything is
rendered.
- """
- model = (
- output_schema.output_schema if isinstance(output_schema, OutputSchema)
else None
- )
- if not (isinstance(model, type) and issubclass(model, BaseModel)):
- return None
- return model
-
-
def _native_response_format(
output_schema: OutputSchema | None,
) -> Dict[str, Any] | None:
@@ -197,7 +182,7 @@ def _native_response_format(
Rendering goes through the shared helper, which raises ``TypeError``
naming the
schema class and the remedy when the model has no JSON Schema.
"""
- model = _native_output_model(output_schema)
+ model = native_output_model(output_schema)
if model is None:
return None
return {
@@ -451,7 +436,7 @@ class WatsonxChatModelConnection(BaseChatModelConnection):
bool
``True`` if ``output_schema`` wraps a ``BaseModel`` subclass.
"""
- return _native_output_model(output_schema) is not None
+ return native_output_model(output_schema) is not None
def chat(
self,
@@ -512,7 +497,7 @@ class WatsonxChatModelConnection(BaseChatModelConnection):
# Feasibility is asked rather than restated, so a caller asking the
same
# question gets the answer this branch acts on.
if self._can_apply_native_structured_output(output_schema, tools,
raw_kwargs):
- native_model = _native_output_model(output_schema)
+ native_model = native_output_model(output_schema)
# A caller reaches the same request field through either channel,
and both
# have already merged into request_params. Only the branch that
sends a
# derived schema may reject the caller's value; every path that
skips it
