Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions packages/gen/gen_ai_hub/orchestration_v2/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,7 @@

# multimodal_items
"ImageDetailLevel", "TextPart", "ImageUrl", "ImagePart", "ContentPart", "ImageItem",
"FileContent", "FilePart", "FileItem",

# response
"CacheCreationTokenDetails", "PromptTokensDetails", "CompletionTokensDetails", "TokenUsage", "GenericModuleResult", "TopLogprob",
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,8 @@
from .llm_model_details import LLMModelDetails
from .message import (SystemMessage, UserMessage, AssistantMessage, ToolChatMessage, DeveloperChatMessage, ChatMessage,
ResponseChatMessage, FunctionCall, MessageToolCall, ReasoningBlock)
from .multimodal_items import ImageDetailLevel, TextPart, ImageUrl, ImagePart, ContentPart, ImageItem
from .multimodal_items import (ImageDetailLevel, TextPart, ImageUrl, ImagePart, ContentPart, ImageItem,
FileContent, FilePart, FileItem)
from .response import (CacheCreationTokenDetails, PromptTokensDetails, CompletionTokensDetails, TokenUsage, GenericModuleResult, TopLogprob,
ChatCompletionTokenLogprob, ChoiceLogprobs, LLMChoice, StreamFunctionObject, StreamToolCall,
StreamDelta, StreamLLMChoice, Citation, LLMModuleResult, StreamLLMModuleResult, ModuleResults,
Expand Down Expand Up @@ -81,6 +82,7 @@

# multimodal_items
"ImageDetailLevel", "TextPart", "ImageUrl", "ImagePart", "ContentPart", "ImageItem",
"FileContent", "FilePart", "FileItem",

# response
"CacheCreationTokenDetails", "PromptTokensDetails", "CompletionTokensDetails", "TokenUsage", "GenericModuleResult", "TopLogprob",
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@

from typing import Optional, Dict

from pydantic import Field
from pydantic import Field, StrictBool

from gen_ai_hub.orchestration_v2.models.base import ABCBaseModel as BaseModel

Expand All @@ -24,11 +24,18 @@ class LLMModelDetails(BaseModel):
timeout: Timeout for the LLM request in seconds. This parameter is currently ignored for Vertex AI models.

max_retries: Maximum number of retries for the LLM request. This parameter is currently ignored for Vertex AI models.

drop_unsupported_params: Controls behaviour when Chat Completions parameters not supported by the
Responses API are present in the request. Only applies to models that use the Responses API.
When True (default): unsupported params are silently dropped and reported in the
X-Orchestration-Dropped-Model-Params response header. When False: any unsupported parameter
causes a 400 error with an actionable message.
"""
name: str
version: Optional[str] = "latest"
params: Optional[Dict] = None
timeout: Optional[int] = Field(default=600, ge=1, le=600)
timeout: Optional[int] = Field(default=600, ge=1, le=1200)
max_retries: Optional[int] = Field(default=2, ge=0, le=5)
drop_unsupported_params: Optional[StrictBool] = Field(default=True)

__all__ = ["LLMModelDetails"]
12 changes: 8 additions & 4 deletions packages/gen/gen_ai_hub/orchestration_v2/models/message.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,9 @@
from pydantic import field_validator, ValidationError

from gen_ai_hub.orchestration_v2.models.base import ABCBaseModel as BaseModel, ResponseBaseModel
from gen_ai_hub.orchestration_v2.models.multimodal_items import ContentPart, ImageItem, TextPart, ImageUrl, ImagePart
from gen_ai_hub.orchestration_v2.models.multimodal_items import (
ContentPart, ImageItem, TextPart, ImageUrl, ImagePart, FileItem, FilePart, FileContent
)


class FunctionCall(ResponseBaseModel):
Expand Down Expand Up @@ -122,10 +124,10 @@ class UserMessage(BaseModel):
Args:
role: The role of the entity sending the message.

content: The message content, which may be plain text or a sequence of text and images.
content: The message content, which may be plain text or a sequence of text, images, or files.
"""
role: Role = Role.USER
content: Union[str, ContentPart, List[Union[str, ContentPart, ImageItem, ImagePart]]]
content: Union[str, ContentPart, List[Union[str, ContentPart, ImageItem, ImagePart, FileItem, FilePart]]]

@field_validator("content", mode="before")
def content_validation(cls, content): # pylint: disable=no-self-argument
Expand All @@ -145,8 +147,10 @@ def content_validation(cls, content): # pylint: disable=no-self-argument
mapped_content.append(TextPart(text=item))
elif isinstance(item, ImageItem):
mapped_content.append(ImagePart(image_url=ImageUrl(url=item.url, detail=item.detail)))
elif isinstance(item, FileItem):
mapped_content.append(FilePart(file=FileContent(file_data=item.file_data, filename=item.filename)))
else:
raise ValidationError("User message content list must contain only str or ImageItem")
raise ValidationError("User message content list must contain only str, ImageItem or FileItem")
return mapped_content


Expand Down
114 changes: 111 additions & 3 deletions packages/gen/gen_ai_hub/orchestration_v2/models/multimodal_items.py
Original file line number Diff line number Diff line change
Expand Up @@ -76,7 +76,36 @@ class ImagePart(BaseModel):
cache_control: Optional[CacheControl] = None


ContentPart = Union[TextPart, ImagePart]
class FileContent(BaseModel):
"""
A data structure holding the data and name for a file.

Args:
file_data: The file content, as a base64 encoded string or a file URL.

filename: The name of the file.
"""
file_data: str
filename: Optional[str] = None


class FilePart(BaseModel):
"""
Represents a file segment within a multimodal content block.

Args:
file: A `FileContent` object containing the file data and optional filename.

type: The type identifier, defaulting to "file".

cache_control: Optional cache control settings for prompt caching.
"""
file: FileContent
type_: Literal["file"] = Field(default="file", alias="type")
cache_control: Optional[CacheControl] = None


ContentPart = Union[TextPart, ImagePart, FilePart]


class ImageItem(BaseModel):
Expand Down Expand Up @@ -149,7 +178,7 @@ def model_dump( # pylint: disable=arguments-differ
warnings: bool | Literal['none', 'warn', 'error'] = True,
fallback: Callable[[Any], Any] | None = None,
serialize_as_any: bool = False,
) -> dict[str, Any]:
) -> dict[str, Any]:
return ImagePart(image_url=ImageUrl(url=self.url, detail=self.detail)).model_dump(
mode=mode,
include=include,
Expand All @@ -165,4 +194,83 @@ def model_dump( # pylint: disable=arguments-differ
serialize_as_any=serialize_as_any,
)

__all__ = ["ImageDetailLevel", "TextPart", "ImageUrl", "ImagePart", "ContentPart", "ImageItem"]

class FileItem(BaseModel):
"""
Represents a file for use in multimodal messages.

Args:
file_data: The file content, specified as either a base64 encoded string,
a data URL, or a file URL.

filename: The name of the file.
"""

file_data: Optional[str] = None
filename: Optional[str] = None

@staticmethod
def from_file(
file_path: str,
mime_type: Optional[str] = None,
filename: Optional[str] = None,
) -> "FileItem":
"""Create a FileItem from a local file.

:param file_path: Path to the file.
:type file_path: str
:param mime_type: Explicit MIME type (e.g., 'application/pdf').
If not provided, the MIME type will be guessed from the file extension.
:type mime_type: Optional[str], optional
:param filename: The name of the file.
:type filename: Optional[str], optional
:raises ValueError: If the MIME type cannot be determined and is not provided.
:return: A FileItem instance with the file data as a data URL.
:rtype: FileItem
"""

mime = mime_type or mimetypes.guess_type(file_path)[0]
if not mime:
raise ValueError(
f"Could not determine MIME type for file: {file_path}. "
"Please provide mime_type explicitly."
)
with open(file_path, "rb") as file:
encoded = base64.b64encode(file.read()).decode("utf-8")
data_url = f"data:{mime};base64,{encoded}"

return FileItem(file_data=data_url, filename=filename)

def model_dump( # pylint: disable=arguments-differ
self,
*,
mode: Literal['json', 'python'] | str = 'python',
include: IncEx | None = None,
exclude: IncEx | None = None,
context: Any | None = None,
by_alias: bool = True,
exclude_unset: bool = False,
exclude_defaults: bool = False,
exclude_none: bool = False,
round_trip: bool = False,
warnings: bool | Literal['none', 'warn', 'error'] = True,
fallback: Callable[[Any], Any] | None = None,
serialize_as_any: bool = False,
) -> dict[str, Any]:
return FilePart(file=FileContent(file_data=self.file_data, filename=self.filename)).model_dump(
mode=mode,
include=include,
exclude=exclude,
context=context,
by_alias=by_alias,
exclude_unset=exclude_unset,
exclude_defaults=exclude_defaults,
exclude_none=exclude_none,
round_trip=round_trip,
warnings=warnings,
fallback=fallback,
serialize_as_any=serialize_as_any,
)

__all__ = ["ImageDetailLevel", "TextPart", "ImageUrl", "ImagePart", "ContentPart", "ImageItem",
"FileContent", "FilePart", "FileItem"]
6 changes: 6 additions & 0 deletions packages/gen/gen_ai_hub/orchestration_v2/models/response.py
Original file line number Diff line number Diff line change
Expand Up @@ -144,12 +144,17 @@ class LLMChoice(ResponseBaseModel):

- 'content_filter' if content was omitted due to a filter enforced by the LLM model provider
or the content filtering module

provider_specific_fields: Provider-specific fields not representable in the OpenAI schema.
When not identical, includes 'finish_reason' with the raw, unmapped
finish reason string from the provider.
"""

index: int
message: ResponseChatMessage
logprobs: Optional[ChoiceLogprobs] = None
finish_reason: str
provider_specific_fields: Optional[dict[str, Any]] = None


class StreamFunctionObject(FunctionCall):
Expand Down Expand Up @@ -197,6 +202,7 @@ class StreamLLMChoice(ResponseBaseModel):
delta: StreamDelta
logprobs: Optional[ChoiceLogprobs] = None
finish_reason: Optional[str] = None
provider_specific_fields: Optional[dict[str, Any]] = None


class Citation(ResponseBaseModel):
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@ class Template(BaseModel):


class PromptTemplatingModuleConfig(BaseModel):
prompt: Template | TemplateRef
prompt: Optional[Template | TemplateRef] = None
model: LLMModelDetails

__all__ = ["Template", "PromptTemplatingModuleConfig"]
8 changes: 8 additions & 0 deletions packages/gen/tests/orchestration_v2/test_config_v2.py
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,14 @@ def test_minimal_config(self):
)
self.assertNotIn("filtering_module_config", json_data["modules"])

def test_prompt_optional(self):
# prompt is no longer required; it is omitted from the payload when not provided
# (messages_history must then be supplied in the request).
prompt_template = PromptTemplatingModuleConfig(model=self.llm)
self.assertIsNone(prompt_template.prompt)
dumped = prompt_template.model_dump()
self.assertNotIn("prompt", dumped)

def test_input_filtering(self):
input_filtering = InputFiltering(filters=[
AzureContentSafetyInputFilterConfig(config=AzureContentSafetyInput(hate=0))
Expand Down
1 change: 1 addition & 0 deletions packages/gen/tests/orchestration_v2/test_flat_import.py
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,7 @@

# multimodal_items
"ImageDetailLevel", "TextPart", "ImageUrl", "ImagePart", "ContentPart", "ImageItem",
"FileContent", "FilePart", "FileItem",

# response
"PromptTokensDetails", "CompletionTokensDetails", "TokenUsage", "GenericModuleResult", "TopLogprob",
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -23,11 +23,27 @@ def test_llm_json_serialization(self):
"version": "v1",
"params": {"temperature": 0.7},
"max_retries": 2,
"timeout": 600
"timeout": 600,
"drop_unsupported_params": True
}
json_data = llm.model_dump()
self.assertEqual(json_data, expected_dict)

def test_llm_timeout_upper_bound(self):
llm = LLMModelDetails(name="gpt-4o-mini", timeout=1200)
self.assertEqual(llm.timeout, 1200)

with self.assertRaises(ValueError):
LLMModelDetails(name="gpt-4o-mini", timeout=1201)

def test_llm_drop_unsupported_params(self):
llm = LLMModelDetails(name="gpt-4o-mini")
self.assertEqual(llm.drop_unsupported_params, True)

llm = LLMModelDetails(name="gpt-4o-mini", drop_unsupported_params=False)
self.assertEqual(llm.drop_unsupported_params, False)
self.assertEqual(llm.model_dump()["drop_unsupported_params"], False)

def test_llm_model_details_with_unsupported_properties(self):
with self.assertRaises(ValueError):
LLMModelDetails(name="gpt-4o-mini", max_retries=10)
Expand Down
33 changes: 33 additions & 0 deletions packages/gen/tests/orchestration_v2/test_message_v2.py
Original file line number Diff line number Diff line change
Expand Up @@ -102,6 +102,23 @@ def test_llm_choice_stores_extra_field(self):
})
self.assertEqual(choice.extra_field, "extra")

def test_llm_choice_provider_specific_fields(self):
choice = LLMChoice.model_validate({
"index": 0,
"message": {"role": "assistant", "content": "ok"},
"finish_reason": "stop",
"provider_specific_fields": {"finish_reason": "end_turn"},
})
self.assertEqual(choice.provider_specific_fields, {"finish_reason": "end_turn"})

def test_llm_choice_provider_specific_fields_optional(self):
choice = LLMChoice.model_validate({
"index": 0,
"message": {"role": "assistant", "content": "ok"},
"finish_reason": "stop",
})
self.assertIsNone(choice.provider_specific_fields)

def test_llm_module_result_stores_extra_field(self):
result = LLMModuleResult.model_validate({
"id": "chatcmpl-abc",
Expand Down Expand Up @@ -189,6 +206,22 @@ def test_stream_llm_choice_stores_extra_field(self):
})
self.assertEqual(choice.extra_field, "extra")

def test_stream_llm_choice_provider_specific_fields(self):
choice = StreamLLMChoice.model_validate({
"index": 0,
"delta": {"content": "hi"},
"finish_reason": "stop",
"provider_specific_fields": {"finish_reason": "end_turn"},
})
self.assertEqual(choice.provider_specific_fields, {"finish_reason": "end_turn"})

def test_stream_llm_choice_provider_specific_fields_optional(self):
choice = StreamLLMChoice.model_validate({
"index": 0,
"delta": {"content": "hi"},
})
self.assertIsNone(choice.provider_specific_fields)

def test_citation_stores_extra_field(self):
citation = Citation.model_validate({
"title": "Source A", "url": "https://example.com", "extra_field": "extra",
Expand Down
Loading
Loading