aiqtoolkit 1.1.0a20250429__py3-none-any.whl

aiq/data_models/__init__.py

@@ -0,0 +1,14 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

aiq/data_models/api_server.py

@@ -0,0 +1,550 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import abc
+import datetime
+import typing
+import uuid
+from abc import abstractmethod
+from enum import Enum
+
+from pydantic import BaseModel
+from pydantic import ConfigDict
+from pydantic import Discriminator
+from pydantic import HttpUrl
+from pydantic import conlist
+from pydantic import field_validator
+from pydantic_core.core_schema import ValidationInfo
+
+from aiq.data_models.interactive import HumanPrompt
+from aiq.utils.type_converter import GlobalTypeConverter
+
+
+class Message(BaseModel):
+    content: str
+    role: str
+
+
+class AIQChatRequest(BaseModel):
+    """
+    AIQChatRequest is a data model that represents a request to the AgentIQ chat API.
+    """
+
+    # Allow extra fields in the model_config to support derived models
+    model_config = ConfigDict(extra="allow")
+
+    messages: typing.Annotated[list[Message], conlist(Message, min_length=1)]
+    model: str | None = None
+    temperature: float | None = None
+    max_tokens: int | None = None
+    top_p: float | None = None
+
+    @staticmethod
+    def from_string(data: str,
+                    *,
+                    model: str | None = None,
+                    temperature: float | None = None,
+                    max_tokens: int | None = None,
+                    top_p: float | None = None) -> "AIQChatRequest":
+
+        return AIQChatRequest(messages=[Message(content=data, role="user")],
+                              model=model,
+                              temperature=temperature,
+                              max_tokens=max_tokens,
+                              top_p=top_p)
+
+
+class AIQChoiceMessage(BaseModel):
+    content: str | None = None
+    role: str | None = None
+
+
+class AIQChoice(BaseModel):
+    model_config = ConfigDict(extra="allow")
+
+    message: AIQChoiceMessage
+    finish_reason: typing.Literal['stop', 'length', 'tool_calls', 'content_filter', 'function_call'] | None = None
+    index: int
+    # logprobs: AIQChoiceLogprobs | None = None
+
+
+class AIQUsage(BaseModel):
+    prompt_tokens: int
+    completion_tokens: int
+    total_tokens: int
+
+
+class AIQResponseSerializable(abc.ABC):
+    """
+    AIQChatResponseSerializable is an abstract class that defines the interface for serializing output for the AgentIQ
+    chat streaming API.
+    """
+
+    @abstractmethod
+    def get_stream_data(self) -> str:
+        pass
+
+
+class AIQResponseBaseModelOutput(BaseModel, AIQResponseSerializable):
+
+    def get_stream_data(self) -> str:
+        return f"data: {self.model_dump_json()}\n\n"
+
+
+class AIQResponseBaseModelIntermediate(BaseModel, AIQResponseSerializable):
+
+    def get_stream_data(self) -> str:
+        return f"intermediate_data: {self.model_dump_json()}\n\n"
+
+
+class AIQChatResponse(AIQResponseBaseModelOutput):
+    """
+    AIQChatResponse is a data model that represents a response from the AgentIQ chat API.
+    """
+
+    # Allow extra fields in the model_config to support derived models
+    model_config = ConfigDict(extra="allow")
+    id: str
+    object: str
+    model: str = ""
+    created: datetime.datetime
+    choices: list[AIQChoice]
+    usage: AIQUsage | None = None
+
+    @staticmethod
+    def from_string(data: str,
+                    *,
+                    id_: str | None = None,
+                    object_: str | None = None,
+                    model: str | None = None,
+                    created: datetime.datetime | None = None,
+                    usage: AIQUsage | None = None) -> "AIQChatResponse":
+
+        if id_ is None:
+            id_ = str(uuid.uuid4())
+        if object_ is None:
+            object_ = "chat.completion"
+        if model is None:
+            model = ""
+        if created is None:
+            created = datetime.datetime.now(datetime.timezone.utc)
+
+        return AIQChatResponse(
+            id=id_,
+            object=object_,
+            model=model,
+            created=created,
+            choices=[AIQChoice(index=0, message=AIQChoiceMessage(content=data), finish_reason="stop")],
+            usage=usage)
+
+
+class AIQChatResponseChunk(AIQResponseBaseModelOutput):
+    """
+    AIQChatResponseChunk is a data model that represents a response chunk from the AgentIQ chat streaming API.
+    """
+
+    # Allow extra fields in the model_config to support derived models
+    model_config = ConfigDict(extra="allow")
+
+    id: str
+    choices: list[AIQChoice]
+    created: datetime.datetime
+    model: str = ""
+    object: str = "chat.completion.chunk"
+
+    @staticmethod
+    def from_string(data: str,
+                    *,
+                    id_: str | None = None,
+                    created: datetime.datetime | None = None,
+                    model: str | None = None,
+                    object_: str | None = None) -> "AIQChatResponseChunk":
+
+        if id_ is None:
+            id_ = str(uuid.uuid4())
+        if created is None:
+            created = datetime.datetime.now(datetime.timezone.utc)
+        if model is None:
+            model = ""
+        if object_ is None:
+            object_ = "chat.completion.chunk"
+
+        return AIQChatResponseChunk(
+            id=id_,
+            choices=[AIQChoice(index=0, message=AIQChoiceMessage(content=data), finish_reason="stop")],
+            created=created,
+            model=model,
+            object=object_)
+
+
+class AIQResponseIntermediateStep(AIQResponseBaseModelIntermediate):
+    """
+    AIQResponseSerializedStep is a data model that represents a serialized step in the AgentIQ chat streaming API.
+    """
+
+    # Allow extra fields in the model_config to support derived models
+    model_config = ConfigDict(extra="allow")
+
+    id: str
+    parent_id: str | None = None
+    type: str = "markdown"
+    name: str
+    payload: str
+
+
+class AIQResponsePayloadOutput(BaseModel, AIQResponseSerializable):
+
+    payload: typing.Any
+
+    def get_stream_data(self) -> str:
+
+        if (isinstance(self.payload, BaseModel)):
+            return f"data: {self.payload.model_dump_json()}\n\n"
+
+        return f"data: {self.payload}\n\n"
+
+
+class AIQGenerateResponse(BaseModel):
+    # Allow extra fields in the model_config to support derived models
+    model_config = ConfigDict(extra="allow")
+
+    # (fixme) define the intermediate step model
+    intermediate_steps: list[tuple] | None = None
+    output: str
+    value: str | None = "default"
+
+
+class UserMessageContentRoleType(str, Enum):
+    USER = "user"
+    ASSISTANT = "assistant"
+
+
+class ChatContentType(str, Enum):
+    """
+    ChatContentType is an Enum that represents the type of Chat content.
+    """
+    TEXT = "text"
+    IMAGE_URL = "image_url"
+    INPUT_AUDIO = "input_audio"
+
+
+class WebSocketMessageType(str, Enum):
+    """
+    WebSocketMessageType is an Enum that represents WebSocket Message types.
+    """
+    USER_MESSAGE = "user_message"
+    RESPONSE_MESSAGE = "system_response_message"
+    INTERMEDIATE_STEP_MESSAGE = "system_intermediate_message"
+    SYSTEM_INTERACTION_MESSAGE = "system_interaction_message"
+    USER_INTERACTION_MESSAGE = "user_interaction_message"
+    ERROR_MESSAGE = "error_message"
+
+
+class WorkflowSchemaType(str, Enum):
+    """
+    WorkflowSchemaType is an Enum that represents Workkflow response types.
+    """
+    GENERATE_STREAM = "generate_stream"
+    CHAT_STREAM = "chat_stream"
+    GENERATE = "generate"
+    CHAT = "chat"
+
+
+class WebSocketMessageStatus(str, Enum):
+    """
+    WebSocketMessageStatus is an Enum that represents the status of a WebSocket message.
+    """
+    IN_PROGRESS = "in_progress"
+    COMPLETE = "complete"
+
+
+class InputAudio(BaseModel):
+    data: str = "default"
+    format: str = "default"
+
+
+class AudioContent(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    type: typing.Literal[ChatContentType.INPUT_AUDIO] = ChatContentType.INPUT_AUDIO
+    input_audio: InputAudio = InputAudio()
+
+
+class ImageUrl(BaseModel):
+    url: HttpUrl = HttpUrl(url="http://default.com")
+
+
+class ImageContent(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    type: typing.Literal[ChatContentType.IMAGE_URL] = ChatContentType.IMAGE_URL
+    image_url: ImageUrl = ImageUrl()
+
+
+class TextContent(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    type: typing.Literal[ChatContentType.TEXT] = ChatContentType.TEXT
+    text: str = "default"
+
+
+class Security(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    api_key: str = "default"
+    token: str = "default"
+
+
+UserContent = typing.Annotated[TextContent | ImageContent | AudioContent, Discriminator("type")]
+
+
+class UserMessages(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    role: UserMessageContentRoleType
+    content: list[UserContent]
+
+
+class UserMessageContent(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+    messages: list[UserMessages]
+
+
+class User(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    name: str = "default"
+    email: str = "default"
+
+
+class ErrorTypes(str, Enum):
+    UNKNOWN_ERROR = "unknown_error"
+    INVALID_MESSAGE = "invalid_message"
+    INVALID_MESSAGE_TYPE = "invalid_message_type"
+    INVALID_USER_MESSAGE_CONTENT = "invalid_user_message_content"
+    INVALID_DATA_CONTENT = "invalid_data_content"
+
+
+class Error(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    code: ErrorTypes = ErrorTypes.UNKNOWN_ERROR
+    message: str = "default"
+    details: str = "default"
+
+
+class WebSocketUserMessage(BaseModel):
+    """
+    For more details, refer to the API documentation:
+    docs/source/developer_guide/websockets.md
+    """
+    # Allow extra fields in the model_config to support derived models
+    model_config = ConfigDict(extra="allow")
+
+    type: typing.Literal[WebSocketMessageType.USER_MESSAGE]
+    schema_type: WorkflowSchemaType
+    id: str = "default"
+    thread_id: str = "default"
+    content: UserMessageContent
+    user: User = User()
+    security: Security = Security()
+    error: Error = Error()
+    schema_version: str = "1.0.0"
+    timestamp: str = str(datetime.datetime.now(datetime.timezone.utc))
+
+
+class WebSocketUserInteractionResponseMessage(BaseModel):
+    """
+    For more details, refer to the API documentation:
+    docs/source/developer_guide/websockets.md
+    """
+    type: typing.Literal[WebSocketMessageType.USER_INTERACTION_MESSAGE]
+    id: str = "default"
+    thread_id: str = "default"
+    content: UserMessageContent
+    user: User = User()
+    security: Security = Security()
+    error: Error = Error()
+    schema_version: str = "1.0.0"
+    timestamp: str = str(datetime.datetime.now(datetime.timezone.utc))
+
+
+class SystemIntermediateStepContent(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+    name: str
+    payload: str
+
+
+class WebSocketSystemIntermediateStepMessage(BaseModel):
+    """
+    For more details, refer to the API documentation:
+    docs/source/developer_guide/websockets.md
+    """
+    # Allow extra fields in the model_config to support derived models
+    model_config = ConfigDict(extra="allow")
+
+    type: typing.Literal[WebSocketMessageType.INTERMEDIATE_STEP_MESSAGE]
+    id: str = "default"
+    thread_id: str | None = "default"
+    parent_id: str = "default"
+    intermediate_parent_id: str | None = "default"
+    update_message_id: str | None = "default"
+    content: SystemIntermediateStepContent
+    status: WebSocketMessageStatus
+    timestamp: str = str(datetime.datetime.now(datetime.timezone.utc))
+
+
+class SystemResponseContent(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    text: str | None = None
+
+
+class WebSocketSystemResponseTokenMessage(BaseModel):
+    """
+    For more details, refer to the API documentation:
+    docs/source/developer_guide/websockets.md
+    """
+    # Allow extra fields in the model_config to support derived models
+    model_config = ConfigDict(extra="allow")
+
+    type: typing.Literal[WebSocketMessageType.RESPONSE_MESSAGE, WebSocketMessageType.ERROR_MESSAGE]
+    id: str | None = "default"
+    thread_id: str | None = "default"
+    parent_id: str = "default"
+    content: SystemResponseContent | Error | AIQGenerateResponse
+    status: WebSocketMessageStatus
+    timestamp: str = str(datetime.datetime.now(datetime.timezone.utc))
+
+    @field_validator("content")
+    @classmethod
+    def validate_content_by_type(cls, value: SystemResponseContent | Error | AIQGenerateResponse, info: ValidationInfo):
+        if info.data.get("type") == WebSocketMessageType.ERROR_MESSAGE and not isinstance(value, Error):
+            raise ValueError(f"Field: content must be 'Error' when type is {WebSocketMessageType.ERROR_MESSAGE}")
+
+        if info.data.get("type") == WebSocketMessageType.RESPONSE_MESSAGE and not isinstance(
+                value, (SystemResponseContent, AIQGenerateResponse)):
+            raise ValueError(
+                f"Field: content must be 'SystemResponseContent' when type is {WebSocketMessageType.RESPONSE_MESSAGE}")
+        return value
+
+
+class WebSocketSystemInteractionMessage(BaseModel):
+    """
+    For more details, refer to the API documentation:
+    docs/source/developer_guide/websockets.md
+    """
+    # Allow extra fields in the model_config to support derived models
+    model_config = ConfigDict(extra="allow")
+
+    type: typing.Literal[
+        WebSocketMessageType.SYSTEM_INTERACTION_MESSAGE] = WebSocketMessageType.SYSTEM_INTERACTION_MESSAGE
+    id: str | None = "default"
+    thread_id: str | None = "default"
+    parent_id: str = "default"
+    content: HumanPrompt
+    status: WebSocketMessageStatus
+    timestamp: str = str(datetime.datetime.now(datetime.timezone.utc))
+
+
+# ======== AIQGenerateResponse Converters ========
+
+
+def _generate_response_to_str(response: AIQGenerateResponse) -> str:
+    return response.output
+
+
+GlobalTypeConverter.register_converter(_generate_response_to_str)
+
+
+def _generate_response_to_chat_response(response: AIQGenerateResponse) -> AIQChatResponse:
+    data = response.output
+
+    # Simulate usage
+    prompt_tokens = 0
+    usage = AIQUsage(prompt_tokens=prompt_tokens,
+                     completion_tokens=len(data.split()),
+                     total_tokens=prompt_tokens + len(data.split()))
+
+    # Build and return the response
+    return AIQChatResponse.from_string(data, usage=usage)
+
+
+GlobalTypeConverter.register_converter(_generate_response_to_chat_response)
+
+
+# ======== AIQChatRequest Converters ========
+def _aiq_chat_request_to_string(data: AIQChatRequest) -> str:
+    return data.messages[-1].content
+
+
+GlobalTypeConverter.register_converter(_aiq_chat_request_to_string)
+
+
+def _string_to_aiq_chat_request(data: str) -> AIQChatRequest:
+    return AIQChatRequest.from_string(data, model="")
+
+
+GlobalTypeConverter.register_converter(_string_to_aiq_chat_request)
+
+
+# ======== AIQChatResponse Converters ========
+def _aiq_chat_response_to_string(data: AIQChatResponse) -> str:
+    return data.choices[0].message.content or ""
+
+
+GlobalTypeConverter.register_converter(_aiq_chat_response_to_string)
+
+
+def _string_to_aiq_chat_response(data: str) -> AIQChatResponse:
+    '''Converts a string to an AIQChatResponse object'''
+
+    # Simulate usage
+    prompt_tokens = 0
+    usage = AIQUsage(prompt_tokens=prompt_tokens,
+                     completion_tokens=len(data.split()),
+                     total_tokens=prompt_tokens + len(data.split()))
+
+    # Build and return the response
+    return AIQChatResponse.from_string(data, usage=usage)
+
+
+GlobalTypeConverter.register_converter(_string_to_aiq_chat_response)
+
+
+def _chat_response_to_chat_response_chunk(data: AIQChatResponse) -> AIQChatResponseChunk:
+
+    return AIQChatResponseChunk(id=data.id, choices=data.choices, created=data.created, model=data.model)
+
+
+GlobalTypeConverter.register_converter(_chat_response_to_chat_response_chunk)
+
+
+# ======== AIQChatResponseChunk Converters ========
+def _aiq_chat_response_chunk_to_string(data: AIQChatResponseChunk) -> str:
+    return data.choices[0].message.content or ""
+
+
+GlobalTypeConverter.register_converter(_aiq_chat_response_chunk_to_string)
+
+
+def _string_to_aiq_chat_response_chunk(data: str) -> AIQChatResponseChunk:
+    '''Converts a string to an AIQChatResponseChunk object'''
+
+    # Build and return the response
+    return AIQChatResponseChunk.from_string(data)
+
+
+GlobalTypeConverter.register_converter(_string_to_aiq_chat_response_chunk)

aiq/data_models/common.py

@@ -0,0 +1,143 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import inspect
+import sys
+import typing
+from hashlib import sha512
+
+from pydantic import AliasChoices
+from pydantic import BaseModel
+from pydantic import Field
+
+_LT = typing.TypeVar("_LT")
+
+
+class HashableBaseModel(BaseModel):
+    """
+    Subclass of a Pydantic BaseModel that is hashable. Use in objects that need to be hashed for caching purposes.
+    """
+
+    def __hash__(self):
+        return int.from_bytes(bytes=sha512(f"{self.__class__.__qualname__}::{self.model_dump_json()}".encode(
+            'utf-8', errors='ignore')).digest(),
+                              byteorder=sys.byteorder)
+
+    def __lt__(self, other):
+        return self.__hash__() < other.__hash__()
+
+    def __eq__(self, other):
+        return self.__hash__() == other.__hash__()
+
+    def __ne__(self, other):
+        return self.__hash__() != other.__hash__()
+
+    def __gt__(self, other):
+        return self.__hash__() > other.__hash__()
+
+    @classmethod
+    def generate_json_schema(cls) -> dict[str, typing.Any]:
+        return cls.model_json_schema()
+
+    @classmethod
+    def write_json_schema(cls, schema_path: str) -> None:
+
+        import json
+
+        schema = cls.generate_json_schema()
+
+        with open(schema_path, "w", encoding="utf-8") as f:
+            json.dump(schema, f, indent=2)
+
+
+def subclass_depth(cls: type) -> int:
+    """
+    Compute a class' subclass depth.
+    """
+    depth = 0
+    while (cls is not object):
+        cls = cls.__base__
+        depth += 1
+    return depth
+
+
+def _get_origin_or_base(cls: type) -> type:
+    """
+    Get the origin of a type or the base class if it is not a generic.
+    """
+    origin = typing.get_origin(cls)
+    if origin is None:
+        return cls
+    return origin
+
+
+class BaseModelRegistryTag:
+
+    pass
+
+
+class TypedBaseModel(BaseModel):
+    """
+    Subclass of Pydantic BaseModel that allows for specifying the object type. Use in Pydantic discriminated unions.
+    """
+
+    type: str = Field(init=False,
+                      serialization_alias="_type",
+                      validation_alias=AliasChoices('type', '_type'),
+                      description="The type of the object",
+                      title="Type",
+                      repr=False)
+
+    full_type: typing.ClassVar[str]
+
+    def __init_subclass__(cls, name: str | None = None):
+        super().__init_subclass__()
+
+        if (name is not None):
+            module = inspect.getmodule(cls)
+
+            assert module is not None, f"Module not found for class {cls} when registering {name}"
+            package_name: str | None = module.__package__
+
+            # If the package name is not set, then we use the module name. Must have some namespace which will be unique
+            if (not package_name):
+                package_name = module.__name__
+
+            full_name = f"{package_name}/{name}"
+
+            type_field = cls.model_fields.get("type")
+            if type_field is not None:
+                type_field.default = name
+            cls.full_type = full_name
+
+    @classmethod
+    def static_type(cls):
+        return cls.model_fields.get("type").default
+
+    @classmethod
+    def static_full_type(cls):
+        return cls.full_type
+
+    @staticmethod
+    def discriminator(v: typing.Any) -> str | None:
+        # If its serialized, then we use the alias
+        if isinstance(v, dict):
+            return v.get("_type", v.get("type"))
+
+        # Otherwise we use the property
+        return getattr(v, "type")
+
+
+TypedBaseModelT = typing.TypeVar("TypedBaseModelT", bound=TypedBaseModel)