{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "05733998",
   "metadata": {},
   "source": [
    "# GenAI messages Python models\n",
    "\n",
    "This file defines Python models for system instructions, input, and output messages.\n",
    "These models are provided for reference only.\n",
    "\n",
    "The model definitions are based on [Pydantic](https://github.com/pydantic/pydantic)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c4ec73b",
   "metadata": {},
   "outputs": [],
   "source": [
    "pip install pydantic~=2.0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b335c0ec",
   "metadata": {
    "vscode": {
     "languageId": "shellscript"
    }
   },
   "source": [
    "## Common Code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5124fe15",
   "metadata": {},
   "outputs": [],
   "source": [
    "from enum import StrEnum\n",
    "import json\n",
    "from typing import Any, List, Literal, Optional, Union\n",
    "from pydantic import BaseModel, Field, RootModel\n",
    "\n",
    "class TextPart(BaseModel):\n",
    "    \"\"\"\n",
    "    Represents text content sent to or received from the model.\n",
    "    \"\"\"\n",
    "    type: Literal['text'] = Field(description=\"The type of the content captured in this part.\")\n",
    "    content: str = Field(description=\"Text content sent to or received from the model.\")\n",
    "\n",
    "    class Config:\n",
    "        extra = \"allow\"\n",
    "\n",
    "class ToolCallRequestPart(BaseModel):\n",
    "    \"\"\"\n",
    "    Represents a tool call requested by the model.\n",
    "    \"\"\"\n",
    "    type: Literal[\"tool_call\"] = Field(description=\"The type of the content captured in this part.\")\n",
    "    id: Optional[str] = Field(default=None, description=\"Unique identifier for the tool call.\")\n",
    "    name: str = Field(description=\"Name of the tool.\")\n",
    "    arguments: Any = Field(default=None, description=\"Arguments for the tool call.\")\n",
    "\n",
    "    class Config:\n",
    "        extra = \"allow\"\n",
    "\n",
    "class ToolCallResponsePart(BaseModel):\n",
    "    \"\"\"\n",
    "    Represents a tool call result sent to the model.\n",
    "    \"\"\"\n",
    "    type: Literal['tool_call_response'] = Field(description=\"The type of the content captured in this part.\")\n",
    "    id: Optional[str] = Field(default=None, description=\"Unique tool call identifier.\")\n",
    "    response: Any = Field(description=\"Tool call response.\")\n",
    "\n",
    "    class Config:\n",
    "        extra = \"allow\"\n",
    "\n",
    "class GenericServerToolCall(BaseModel):\n",
    "    \"\"\"\n",
    "    Represents an arbitrary server tool call with any type and properties.\n",
    "    This allows for extensibility with custom server tool types.\n",
    "    \"\"\"\n",
    "    type: str = Field(description=\"Type identifier for the server tool call.\")\n",
    "\n",
    "    class Config:\n",
    "        extra = \"allow\"\n",
    "\n",
    "class GenericServerToolCallResponse(BaseModel):\n",
    "    \"\"\"\n",
    "    Represents an arbitrary server tool call response with any type and properties.\n",
    "    This allows for extensibility with custom server tool response types.\n",
    "    \"\"\"\n",
    "    type: str = Field(description=\"Type identifier for the server tool call response.\")\n",
    "\n",
    "    class Config:\n",
    "        extra = \"allow\"\n",
    "\n",
    "class ServerToolCallPart(BaseModel):\n",
    "    \"\"\"\n",
    "    Represents a server-side tool call invocation. Server tool calls are executed by the model\n",
    "    provider on the server side rather than by the client application. Provider-specific tools\n",
    "    (e.g., code_interpreter, web_search) can have well-defined schemas defined by the respective providers.\n",
    "    \"\"\"\n",
    "    type: Literal['server_tool_call'] = Field(description=\"The type of the content captured in this part.\")\n",
    "    id: Optional[str] = Field(default=None, description=\"Unique identifier for the server tool call.\")\n",
    "    name: str = Field(description=\"Name of the server tool.\")\n",
    "    server_tool_call: GenericServerToolCall = Field(\n",
    "        description=\"Polymorphic server tool call details with type discriminator. The structure varies based on the tool type.\"\n",
    "    )\n",
    "\n",
    "    class Config:\n",
    "        extra = \"allow\"\n",
    "\n",
    "class ServerToolCallResponsePart(BaseModel):\n",
    "    \"\"\"\n",
    "    Represents a server-side tool call response. Contains the outcome and details of a server tool\n",
    "    execution. Provider-specific tools (e.g., code_interpreter, web_search) can have well-defined\n",
    "    response schemas defined by the respective providers.\n",
    "    \"\"\"\n",
    "    type: Literal['server_tool_call_response'] = Field(description=\"The type of the content captured in this part.\")\n",
    "    id: Optional[str] = Field(default=None, description=\"Unique server tool call identifier matching the original call.\")\n",
    "    server_tool_call_response: GenericServerToolCallResponse = Field(\n",
    "        description=\"Polymorphic server tool call response with type discriminator. The structure varies based on the tool type.\"\n",
    "    )\n",
    "\n",
    "    class Config:\n",
    "        extra = \"allow\"\n",
    "\n",
    "\n",
    "class Modality(StrEnum):\n",
    "    IMAGE = \"image\"\n",
    "    VIDEO = \"video\"\n",
    "    AUDIO = \"audio\"\n",
    "\n",
    "\n",
    "class ReasoningPart(BaseModel):\n",
    "    \"\"\"\n",
    "    Represents reasoning/thinking content received from the model.\n",
    "    \"\"\"\n",
    "    type: Literal['reasoning'] = Field(description=\"The type of the content captured in this part.\")\n",
    "    content: str = Field(description=\"Reasoning/thinking content received from the model.\")\n",
    "\n",
    "    class Config:\n",
    "        extra = \"allow\"\n",
    "\n",
    "class BlobPart(BaseModel):\n",
    "    \"\"\"Represents blob binary data sent inline to the model\"\"\"\n",
    "    type: Literal[\"blob\"] = Field(description=\"The type of the content captured in this part.\")\n",
    "    mime_type: Optional[str] = Field(default=None, description=\"The IANA MIME type of the attached data.\")\n",
    "    modality: Union[Modality, str] = Field(\n",
    "        description=\"The general modality of the data if it is known. Instrumentations SHOULD also set the mimeType field if the specific type is known.\"\n",
    "    )\n",
    "    content: bytes = Field(description=\"Raw bytes of the attached data. This field SHOULD be encoded as a base64 string when serialized to JSON.\")\n",
    "\n",
    "class FilePart(BaseModel):\n",
    "    \"\"\"Represents an external referenced file sent to the model by file id\"\"\"\n",
    "    type: Literal[\"file\"] = Field(description=\"The type of the content captured in this part.\")\n",
    "    mime_type: Optional[str] = Field(default=None, description=\"The IANA MIME type of the attached data.\")\n",
    "    modality: Union[Modality, str] = Field(\n",
    "        description=\"The general modality of the data if it is known. Instrumentations SHOULD also set the mimeType field if the specific type is known.\"\n",
    "    )\n",
    "    file_id: str = Field(description=\"An identifier referencing a file that was pre-uploaded to the provider.\")\n",
    "\n",
    "    class Config:\n",
    "        extra = \"allow\"\n",
    "\n",
    "class UriPart(BaseModel):\n",
    "    \"\"\"Represents an external referenced file sent to the model by URI\"\"\"\n",
    "    type: Literal[\"uri\"] = Field(description=\"The type of the content captured in this part.\")\n",
    "    mime_type: Optional[str] = Field(default=None, description=\"The IANA MIME type of the attached data.\")\n",
    "    modality: Union[Modality, str] = Field(\n",
    "        description=\"The general modality of the data if it is known. Instrumentations SHOULD also set the mimeType field if the specific type is known.\"\n",
    "    )\n",
    "    uri: str = Field(\n",
    "        description=(\n",
    "            \"A URI referencing attached data. It should not be a base64 data URL, \"\n",
    "            \"which should use the `blob` part instead. The URI may use a scheme known to \"\n",
    "            \"the provider api (e.g. `gs://bucket/object.png`), or be a publicly accessible location.\"\n",
    "        )\n",
    "    )\n",
    "\n",
    "    class Config:\n",
    "        extra = \"allow\"\n",
    "\n",
    "class GenericPart(BaseModel):\n",
    "    \"\"\"\n",
    "    Represents an arbitrary message part with any type and properties.\n",
    "    This allows for extensibility with custom message part types.\n",
    "    \"\"\"\n",
    "    type: str = Field(description=\"The type of the content captured in this part.\")\n",
    "\n",
    "    class Config:\n",
    "        extra = \"allow\"\n",
    "\n",
    "# This Union without discriminator will generate anyOf in JSON schema\n",
    "MessagePart = Union[\n",
    "    TextPart,\n",
    "    ToolCallRequestPart,\n",
    "    ToolCallResponsePart,\n",
    "    ServerToolCallPart,\n",
    "    ServerToolCallResponsePart,\n",
    "    BlobPart,\n",
    "    FilePart,\n",
    "    UriPart,\n",
    "    ReasoningPart,\n",
    "    GenericPart,  # Catch-all for any other type\n",
    "    # Add other message part types here as needed,\n",
    "    # e.g. structured output, hosted tool call, etc.\n",
    "]\n",
    "\n",
    "class Role(StrEnum):\n",
    "    SYSTEM = \"system\"\n",
    "    USER = \"user\"\n",
    "    ASSISTANT = \"assistant\"\n",
    "    TOOL = \"tool\"\n",
    "\n",
    "class ChatMessage(BaseModel):\n",
    "    role: Union[Role, str] = Field(\n",
    "        description=\"Role of the entity that created the message.\")\n",
    "    parts: List[MessagePart] = Field(\n",
    "        description=\"List of message parts that make up the message content.\")\n",
    "    name: Optional[str] = Field(default=None, description=\"The name of the participant.\")\n",
    "\n",
    "    class Config:\n",
    "        extra = \"allow\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd7b678f",
   "metadata": {},
   "source": [
    "\n",
    "## `gen_ai.input.messages` model\n",
    "\n",
    "Corresponding attribute: [`gen_ai.input.messages`](/docs/registry/attributes/gen-ai.md#gen-ai-input-messages).\n",
    "JSON schema: [`gen_ai.input.messages.json`](../gen-ai-input-messages.json)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0133d3a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "class InputMessages(RootModel[List[ChatMessage]]):\n",
    "    \"\"\"\n",
    "    Represents the list of input messages sent to the model.\n",
    "    \"\"\"\n",
    "    pass\n",
    "\n",
    "# Print the JSON schema for the InputMessages model\n",
    "with open(\"../gen-ai-input-messages.json\", \"w\") as file:\n",
    "    print(json.dumps(InputMessages.model_json_schema(), indent=4), file=file)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "033e5960",
   "metadata": {},
   "source": [
    "## `gen_ai.output.messages` model\n",
    "\n",
    "Corresponding attribute: [`gen_ai.output.messages`](/docs/registry/attributes/gen-ai.md#gen-ai-output-messages).\n",
    "JSON schema: [`gen_ai-output-messages.json`](../gen-ai-output-messages.json)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12a12112",
   "metadata": {},
   "outputs": [],
   "source": [
    "class FinishReason(StrEnum):\n",
    "    \"\"\"\n",
    "    Represents the reason for finishing the generation.\n",
    "    \"\"\"\n",
    "\n",
    "    STOP = \"stop\"\n",
    "    LENGTH = \"length\"\n",
    "    CONTENT_FILTER = \"content_filter\"\n",
    "    TOOL_CALL = \"tool_call\"\n",
    "    ERROR = \"error\"\n",
    "\n",
    "class OutputMessage(ChatMessage):\n",
    "    \"\"\"\n",
    "    Represents an output message generated by the model or agent. The output message captures\n",
    "    specific response (choice, candidate).\n",
    "    \"\"\"\n",
    "    finish_reason: Union[FinishReason, str] = Field(description=\"Reason for finishing the generation.\")\n",
    "\n",
    "class OutputMessages(RootModel[List[OutputMessage]]):\n",
    "    \"\"\"\n",
    "    Represents the list of output messages generated by the model or agent.\n",
    "    \"\"\"\n",
    "    pass\n",
    "\n",
    "# Print the JSON schema for the OutputMessages model\n",
    "with open(\"../gen-ai-output-messages.json\", \"w\") as file:\n",
    "    print(json.dumps(OutputMessages.model_json_schema(), indent=4), file=file)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c035888",
   "metadata": {},
   "source": [
    "## `gen_ai.system_instructions` model\n",
    "\n",
    "Corresponding attribute: [`gen_ai.system_instructions`](/docs/registry/attributes/gen-ai.md#gen-ai-system-instructions).\n",
    "JSON schema: [`gen_ai-system-instructions.json`](../gen-ai-system-instructions.json)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "faa33b95",
   "metadata": {},
   "outputs": [],
   "source": [
    "class SystemInstructions(RootModel[List[MessagePart]]):\n",
    "    \"\"\"\n",
    "    Represents the list of input messages sent to the model.\n",
    "    \"\"\"\n",
    "    pass\n",
    "\n",
    "# Print the JSON schema for the SystemInstructions model\n",
    "with open(\"../gen-ai-system-instructions.json\", \"w\") as file:\n",
    "    print(json.dumps(SystemInstructions.model_json_schema(), indent=4), file=file)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f019c33a",
   "metadata": {},
   "source": [
    "## `gen_ai.tool.definitions` model\n",
    "\n",
    "Corresponding attribute: [`gen_ai.tool.definitions`](/docs/registry/attributes/gen-ai.md#gen-ai-tool-definitions).\n",
    "JSON schema: [`gen_ai-tool-definitions.json`](../gen-ai-tool-definitions.json)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a9e84726",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Annotated\n",
    "\n",
    "from pydantic import GetCoreSchemaHandler, GetJsonSchemaHandler\n",
    "from pydantic_core import core_schema\n",
    "\n",
    "\n",
    "class JsonSchemaDraft7:\n",
    "    \"\"\"Metadata for Pydantic: exported JSON Schema references draft-07; core validation stays permissive.\"\"\"\n",
    "\n",
    "    @classmethod\n",
    "    def __get_pydantic_core_schema__(cls, source_type: Any, handler: GetCoreSchemaHandler):\n",
    "        return core_schema.any_schema()\n",
    "\n",
    "    @classmethod\n",
    "    def __get_pydantic_json_schema__(cls, core_schema_obj, handler: GetJsonSchemaHandler):\n",
    "        # Single-branch anyOf avoids a top-level $ref so Pydantic does not resolve the external meta-schema URL as #/defs.\n",
    "        return {\"anyOf\": [{\"$ref\": \"http://json-schema.org/draft-07/schema#\"}]}\n",
    "\n",
    "\n",
    "# Runtime values are plain dicts (JSON object); type checkers see dict[str, Any]. JsonSchemaDraft7 only customizes JSON Schema export.\n",
    "JsonSchemaDraft7Dict = Annotated[dict[str, Any], JsonSchemaDraft7]\n",
    "\n",
    "\n",
    "class GenericToolDefinition(BaseModel):\n",
    "    \"\"\"\n",
    "    Represents a tool definition in any form.\n",
    "    \"\"\"\n",
    "    type: str = Field(description=\"The type of the tool.\")\n",
    "    name: str = Field(description=\"The name of the tool.\")\n",
    "\n",
    "    class Config:\n",
    "        extra = \"allow\"\n",
    "\n",
    "class FunctionToolDefinition(GenericToolDefinition):\n",
    "    \"\"\"\n",
    "    Represents a tool definition in the form of a function.\n",
    "    \"\"\"\n",
    "    type: Literal[\"function\"] = Field(description=\"The type of the tool.\")\n",
    "    description: Optional[str] = Field(\n",
    "        default=None,\n",
    "        description=(\n",
    "            \"The description of the tool. \"\n",
    "            \"Since this attribute could be large, it's NOT RECOMMENDED to be populated by default. \"\n",
    "            \"Instrumentations MAY provide a way to enable populating this property.\"\n",
    "        )\n",
    "    )\n",
    "    parameters: Optional[JsonSchemaDraft7Dict] = Field(\n",
    "        default=None,\n",
    "        description=(\n",
    "            \"JSON Schema document describing the parameters accepted by the tool. \"\n",
    "            \"The value MUST conform to JSON Schema draft-07. \"\n",
    "            \"Since this attribute could be large, it's NOT RECOMMENDED to be populated by default. \"\n",
    "            \"Instrumentations MAY provide a way to enable populating this property.\"\n",
    "        )\n",
    "    )\n",
    "\n",
    "    class Config:\n",
    "        extra = \"allow\"\n",
    "\n",
    "ToolDefinition = Union[\n",
    "    FunctionToolDefinition,\n",
    "    GenericToolDefinition,  # Catch-all for any other type\n",
    "    # Add other tool definition types here as needed,\n",
    "    # e.g. file search, code interpreter, etc\n",
    "]\n",
    "\n",
    "class ToolDefinitions(RootModel[List[ToolDefinition]]):\n",
    "    \"\"\"\n",
    "    Represents the list of tool definitions available to the GenAI agent or model.\n",
    "    \"\"\"\n",
    "    pass\n",
    "\n",
    "# Print the JSON schema for the ToolDefinitions model\n",
    "with open(\"../gen-ai-tool-definitions.json\", \"w\") as file:\n",
    "    print(json.dumps(ToolDefinitions.model_json_schema(), indent=4), file=file)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "35eecdfb82a2a76",
   "metadata": {},
   "source": [
    "## `gen_ai.retrieval.documents` model\n",
    "\n",
    "Corresponding attribute: [`gen_ai.retrieval.documents`](/docs/registry/attributes/gen-ai.md#gen-ai-retrieval-documents).\n",
    "JSON schema: [`gen-ai-retrieval-documents.json`](../gen-ai-retrieval-documents.json)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "418b6042120d325e",
   "metadata": {},
   "outputs": [],
   "source": [
    "class RetrievalDocument(BaseModel):\n",
    "    \"\"\"\n",
    "    Represents a single document retrieved from a vector database or search system.\n",
    "    \"\"\"\n",
    "    id: str = Field(description=\"A unique identifier for the document.\")\n",
    "    score: float = Field(description=\"The relevance score of the document.\")\n",
    "\n",
    "    class Config:\n",
    "        extra = \"allow\"  # Allows additional properties like content, metadata, title, uri, etc.\n",
    "\n",
    "class RetrievalDocuments(RootModel[List[RetrievalDocument]]):\n",
    "    \"\"\"\n",
    "    Represents the list of documents retrieved from a vector database or search system.\n",
    "    \"\"\"\n",
    "    pass\n",
    "\n",
    "# Print the JSON schema for the RetrievalDocuments model\n",
    "with open(\"../gen-ai-retrieval-documents.json\", \"w\") as file:\n",
    "    print(json.dumps(RetrievalDocuments.model_json_schema(), indent=4), file=file)"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}