gen_ai_models.py

gen_ai_models.py

from enum import Enum
import json
from typing import Annotated, Any, List, Literal, Optional, Union
from pydantic import BaseModel, Field, RootModel

class TextPart(BaseModel):
    """
    Describes text content sent to or received from the model.
    """
    type: Literal['text']
    content: str = Field(description="Text content provided to or received from the model.")
    class Config:
        extra = "allow"

class ToolCallPart(BaseModel):
    """
    Describes a tool call requested by the model.
    """
    type: Literal["tool_call"] = Field(description="Content type.")
    id: str = Field(description="Unique identifier for the tool call.")
    name: str = Field(description="Name of the tool.")
    arguments: Optional[Any] = Field(description="Arguments for the tool call.")
    class Config:
        extra = "allow"

class ToolCallResponsePart(BaseModel):
    """
    Describes a result of tool call sent to the model.
    """
    type: Literal['tool_call_response'] = Field(description="Message type.")
    id: str = Field(description="Unique tool call identifier.")
    result: Optional[Any] = Field(description="Result of the tool call.")
    class Config:
        extra = "allow"

MessagePart = Annotated[
    Union[
        TextPart,
        ToolCallPart,
        ToolCallResponsePart,
        # Add other message part types here as needed,
        # e.g. image URL, image blob, audio URL, structured output, hosted tool call, etc.
    ],
    Field(discriminator='type'),
]

class ChatMessage(BaseModel):
    role: str
    parts: List[MessagePart]
    class Config:
        extra = "allow"

class InputMessages(RootModel[List[ChatMessage]]):
    """
    Describes input messages sent to the model.
    """
    pass

class FinishReason(str, Enum):
    """
    Describes the reason for finishing the generation.
    """

    STOP = "stop"
    LENGTH = "length"
    CONTENT_FILTER = "content_filter"
    TOOL_CALL = "tool_call"
    ERROR = "error"

class OutputMessage(ChatMessage):
    """
    Describes generated output.
    """
    finish_reason: Union[FinishReason, str] = Field(description="Reason for finishing.")

class OutputMessages(RootModel[List[OutputMessage]]):
    """
    Describes the output messages generated by the model.
    """
    pass

class SystemInstructionMessage(BaseModel):
    """
    Describes system instructions.
    """
    role: str
    message: TextPart = Field(description="System instruction message.")
    class Config:
        extra = "allow"

def strip_titles(obj):
    if isinstance(obj, dict):
        # Build a new dict, skipping keys named "title" or "Title"
        return {
            k: strip_titles(v)
            for k, v in obj.items()
            if k.lower() != "title"
        }
    if isinstance(obj, list):
        return [strip_titles(item) for item in obj]
    # scalars (str, int, None, …) – return unchanged
    return obj

print(json.dumps(strip_titles(OutputMessages.model_json_schema()), indent=4))

My proposal, in line with open-telemetry/semantic-conventions#1913

import json
from typing import Annotated, Any, List, Literal, Optional, Union

from pydantic import BaseModel, Field


class TextPart(BaseModel):
    type: Literal['text']
    content: str


# Can adjust the names to e.g. FunctionCallPart
# if we want to allow for other types of tool calls in the future

class ToolCallPart(BaseModel):
    type: Literal['tool_call']
    id: str
    name: str
    arguments: Optional[Any]


class ToolCallResponsePart(BaseModel):
    type: Literal['tool_call_response']
    id: str
    name: Optional[str]
    result: Optional[Any]


MessagePart = Annotated[
    Union[
        TextPart,
        ToolCallPart,
        ToolCallResponsePart,
        # Add other message part types here as needed,
        # e.g. image URL, image blob, audio URL, structured output, hosted tool call, etc.
    ],
    Field(discriminator='type'),
]


class ChatMessage(BaseModel):
    role: str
    parts: List[MessagePart]


class InputMessages(BaseModel):
    messages: List[ChatMessage]


class Choice(BaseModel):
    index: int
    finish_reason: str
    message: ChatMessage

Then instead of:

AssistantMessage(
    role='assistant',
    tool_calls=[
        ToolCall(
            id='1',
            type='function',
            function=FunctionCall(
                name='example_tool',
                arguments={'key': 'value'}
            )
        )
    ]
)

you would have:

ChatMessage(
    role='assistant',
    parts=[
        ToolCallPart(
            type='tool_call',
            id='1',
            name='example_tool',
            arguments={'key': 'value'},
        )
    ],
)