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"]
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
57 changes: 56 additions & 1 deletion packages/gen/tests/orchestration_v2/test_template_v2.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@
UserMessage,
AssistantMessage,
)
from gen_ai_hub.orchestration_v2.models.multimodal_items import ImageItem, ImageDetailLevel
from gen_ai_hub.orchestration_v2.models.multimodal_items import ImageItem, ImageDetailLevel, FileItem
from gen_ai_hub.orchestration_v2.models.response_format import (
ResponseFormatType,
ResponseFormatText,
Expand Down Expand Up @@ -249,3 +249,58 @@ def test_image_item_from_file(self):
self.assertEqual(decoded, original)
finally:
os.unlink(tmp_path)

def test_user_message_with_string_and_file(self):
file = FileItem(file_data="https://example.com/1.pdf", filename="doc.pdf")
msg = UserMessage(content=["Summarize this file:", file])
expected = {
"role": Role.USER,
"content": [
{"type": "text", "text": "Summarize this file:"},
{"type": "file", "file": {
"file_data": "https://example.com/1.pdf",
"filename": "doc.pdf"
}}
]
}
self.assertEqual(msg.model_dump(), expected)

def test_user_message_with_file_dict(self):
msg = UserMessage(content=[{"type": "file", "file": {"file_data": "data:x", "filename": "a.pdf"}}])
expected = {
"role": Role.USER,
"content": [
{"type": "file", "file": {"file_data": "data:x", "filename": "a.pdf"}}
]
}
self.assertEqual(msg.model_dump(), expected)

def test_user_message_with_file_without_filename(self):
file = FileItem(file_data="https://example.com/1.pdf")
msg = UserMessage(content=[file])
expected = {
"role": Role.USER,
"content": [
{"type": "file", "file": {"file_data": "https://example.com/1.pdf"}}
]
}
self.assertEqual(msg.model_dump(), expected)

def test_file_item_from_file(self):
with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp:
tmp.write(b"not a real pdf")
tmp_path = tmp.name

try:
item = FileItem.from_file(tmp_path)
self.assertTrue(item.file_data.startswith("data:application/pdf;base64,"))
self.assertIsNone(item.filename)
# Check that the base64 part decodes to the original file content
encoded = item.file_data.split(",", 1)[1]
with open(tmp_path, "rb") as f:
original = f.read()
decoded = base64.b64decode(encoded)

self.assertEqual(decoded, original)
finally:
os.unlink(tmp_path)
Loading