kagent-adk 0.6.7__py3-none-any.whl → 0.6.9__py3-none-any.whl

kagent/adk/converters/part_converter.py

@@ -0,0 +1,206 @@
+# Copyright 2025 Google LLC
+#
+# 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.
+
+"""
+module containing utilities for conversion betwen A2A Part and Google GenAI Part
+"""
+
+from __future__ import annotations
+
+import base64
+import json
+import logging
+from typing import Optional
+
+from a2a import types as a2a_types
+from google.genai import types as genai_types
+
+from kagent.core.a2a import (
+    A2A_DATA_PART_METADATA_TYPE_CODE_EXECUTION_RESULT,
+    A2A_DATA_PART_METADATA_TYPE_EXECUTABLE_CODE,
+    A2A_DATA_PART_METADATA_TYPE_FUNCTION_CALL,
+    A2A_DATA_PART_METADATA_TYPE_FUNCTION_RESPONSE,
+    A2A_DATA_PART_METADATA_TYPE_KEY,
+    get_kagent_metadata_key,
+)
+
+logger = logging.getLogger("kagent_adk." + __name__)
+
+
+def convert_a2a_part_to_genai_part(
+    a2a_part: a2a_types.Part,
+) -> Optional[genai_types.Part]:
+    """Convert an A2A Part to a Google GenAI Part."""
+    part = a2a_part.root
+    if isinstance(part, a2a_types.TextPart):
+        return genai_types.Part(text=part.text)
+
+    if isinstance(part, a2a_types.FilePart):
+        if isinstance(part.file, a2a_types.FileWithUri):
+            return genai_types.Part(
+                file_data=genai_types.FileData(file_uri=part.file.uri, mime_type=part.file.mime_type)
+            )
+
+        elif isinstance(part.file, a2a_types.FileWithBytes):
+            return genai_types.Part(
+                inline_data=genai_types.Blob(
+                    data=base64.b64decode(part.file.bytes),
+                    mime_type=part.file.mime_type,
+                )
+            )
+        else:
+            logger.warning(
+                "Cannot convert unsupported file type: %s for A2A part: %s",
+                type(part.file),
+                a2a_part,
+            )
+            return None
+
+    if isinstance(part, a2a_types.DataPart):
+        # Conver the Data Part to funcall and function reponse.
+        # This is mainly for converting human in the loop and auth request and
+        # response.
+        # TODO once A2A defined how to suervice such information, migrate below
+        # logic accordinlgy
+        if part.metadata and get_kagent_metadata_key(A2A_DATA_PART_METADATA_TYPE_KEY) in part.metadata:
+            if (
+                part.metadata[get_kagent_metadata_key(A2A_DATA_PART_METADATA_TYPE_KEY)]
+                == A2A_DATA_PART_METADATA_TYPE_FUNCTION_CALL
+            ):
+                return genai_types.Part(function_call=genai_types.FunctionCall.model_validate(part.data, by_alias=True))
+            if (
+                part.metadata[get_kagent_metadata_key(A2A_DATA_PART_METADATA_TYPE_KEY)]
+                == A2A_DATA_PART_METADATA_TYPE_FUNCTION_RESPONSE
+            ):
+                return genai_types.Part(
+                    function_response=genai_types.FunctionResponse.model_validate(part.data, by_alias=True)
+                )
+            if (
+                part.metadata[get_kagent_metadata_key(A2A_DATA_PART_METADATA_TYPE_KEY)]
+                == A2A_DATA_PART_METADATA_TYPE_CODE_EXECUTION_RESULT
+            ):
+                return genai_types.Part(
+                    code_execution_result=genai_types.CodeExecutionResult.model_validate(part.data, by_alias=True)
+                )
+            if (
+                part.metadata[get_kagent_metadata_key(A2A_DATA_PART_METADATA_TYPE_KEY)]
+                == A2A_DATA_PART_METADATA_TYPE_EXECUTABLE_CODE
+            ):
+                return genai_types.Part(
+                    executable_code=genai_types.ExecutableCode.model_validate(part.data, by_alias=True)
+                )
+        return genai_types.Part(text=json.dumps(part.data))
+
+    logger.warning(
+        "Cannot convert unsupported part type: %s for A2A part: %s",
+        type(part),
+        a2a_part,
+    )
+    return None
+
+
+def convert_genai_part_to_a2a_part(
+    part: genai_types.Part,
+) -> Optional[a2a_types.Part]:
+    """Convert a Google GenAI Part to an A2A Part."""
+
+    if part.text:
+        a2a_part = a2a_types.TextPart(text=part.text)
+        if part.thought is not None:
+            a2a_part.metadata = {get_kagent_metadata_key("thought"): part.thought}
+        return a2a_types.Part(root=a2a_part)
+
+    if part.file_data:
+        return a2a_types.Part(
+            root=a2a_types.FilePart(
+                file=a2a_types.FileWithUri(
+                    uri=part.file_data.file_uri,
+                    mime_type=part.file_data.mime_type,
+                )
+            )
+        )
+
+    if part.inline_data:
+        a2a_part = a2a_types.FilePart(
+            file=a2a_types.FileWithBytes(
+                bytes=base64.b64encode(part.inline_data.data).decode("utf-8"),
+                mime_type=part.inline_data.mime_type,
+            )
+        )
+
+        if part.video_metadata:
+            a2a_part.metadata = {
+                get_kagent_metadata_key("video_metadata"): part.video_metadata.model_dump(
+                    by_alias=True, exclude_none=True
+                )
+            }
+
+        return a2a_types.Part(root=a2a_part)
+
+    # Conver the funcall and function reponse to A2A DataPart.
+    # This is mainly for converting human in the loop and auth request and
+    # response.
+    # TODO once A2A defined how to suervice such information, migrate below
+    # logic accordinlgy
+    if part.function_call:
+        return a2a_types.Part(
+            root=a2a_types.DataPart(
+                data=part.function_call.model_dump(by_alias=True, exclude_none=True),
+                metadata={
+                    get_kagent_metadata_key(A2A_DATA_PART_METADATA_TYPE_KEY): A2A_DATA_PART_METADATA_TYPE_FUNCTION_CALL
+                },
+            )
+        )
+
+    if part.function_response:
+        return a2a_types.Part(
+            root=a2a_types.DataPart(
+                data=part.function_response.model_dump(by_alias=True, exclude_none=True),
+                metadata={
+                    get_kagent_metadata_key(
+                        A2A_DATA_PART_METADATA_TYPE_KEY
+                    ): A2A_DATA_PART_METADATA_TYPE_FUNCTION_RESPONSE
+                },
+            )
+        )
+
+    if part.code_execution_result:
+        return a2a_types.Part(
+            root=a2a_types.DataPart(
+                data=part.code_execution_result.model_dump(by_alias=True, exclude_none=True),
+                metadata={
+                    get_kagent_metadata_key(
+                        A2A_DATA_PART_METADATA_TYPE_KEY
+                    ): A2A_DATA_PART_METADATA_TYPE_CODE_EXECUTION_RESULT
+                },
+            )
+        )
+
+    if part.executable_code:
+        return a2a_types.Part(
+            root=a2a_types.DataPart(
+                data=part.executable_code.model_dump(by_alias=True, exclude_none=True),
+                metadata={
+                    get_kagent_metadata_key(
+                        A2A_DATA_PART_METADATA_TYPE_KEY
+                    ): A2A_DATA_PART_METADATA_TYPE_EXECUTABLE_CODE
+                },
+            )
+        )
+
+    logger.warning(
+        "Cannot convert unsupported part for Google GenAI part: %s",
+        part,
+    )
+    return None

kagent/adk/converters/request_converter.py

@@ -0,0 +1,33 @@
+from typing import Any
+
+from a2a.server.agent_execution import RequestContext
+from google.adk.runners import RunConfig
+from google.genai import types as genai_types
+
+from .part_converter import convert_a2a_part_to_genai_part
+
+
+def _get_user_id(request: RequestContext) -> str:
+    # Get user from call context if available (auth is enabled on a2a server)
+    if request.call_context and request.call_context.user and request.call_context.user.user_name:
+        return request.call_context.user.user_name
+
+    # Get user from context id
+    return f"A2A_USER_{request.context_id}"
+
+
+def convert_a2a_request_to_adk_run_args(
+    request: RequestContext,
+) -> dict[str, Any]:
+    if not request.message:
+        raise ValueError("Request message cannot be None")
+
+    return {
+        "user_id": _get_user_id(request),
+        "session_id": request.context_id,
+        "new_message": genai_types.Content(
+            role="user",
+            parts=[convert_a2a_part_to_genai_part(part) for part in request.message.parts],
+        ),
+        "run_config": RunConfig(),
+    }

kagent/adk/models/__init__.py

@@ -0,0 +1,3 @@
+from ._openai import AzureOpenAI, OpenAI
+
+__all__ = ["OpenAI", "AzureOpenAI"]

kagent/adk/models/_openai.py

@@ -0,0 +1,393 @@
+from __future__ import annotations
+
+import base64
+import json
+import os
+from functools import cached_property
+from typing import TYPE_CHECKING, Any, AsyncGenerator, Iterable, Literal, Optional
+
+from google.adk.models import BaseLlm
+from google.adk.models.llm_response import LlmResponse
+from google.genai import types
+from openai import AsyncAzureOpenAI, AsyncOpenAI
+from openai.types.chat import (
+    ChatCompletion,
+    ChatCompletionAssistantMessageParam,
+    ChatCompletionContentPartImageParam,
+    ChatCompletionContentPartTextParam,
+    ChatCompletionMessageParam,
+    ChatCompletionSystemMessageParam,
+    ChatCompletionToolMessageParam,
+    ChatCompletionToolParam,
+    ChatCompletionUserMessageParam,
+)
+from openai.types.chat.chat_completion_message_tool_call_param import (
+    ChatCompletionMessageToolCallParam,
+)
+from openai.types.chat.chat_completion_message_tool_call_param import (
+    Function as ToolCallFunction,
+)
+from openai.types.shared_params import FunctionDefinition
+from pydantic import Field
+
+if TYPE_CHECKING:
+    from google.adk.models.llm_request import LlmRequest
+
+
+def _convert_role_to_openai(role: Optional[str]) -> str:
+    """Convert google.genai role to OpenAI role."""
+    if role in ["model", "assistant"]:
+        return "assistant"
+    elif role == "system":
+        return "system"
+    else:
+        return "user"
+
+
+def _convert_content_to_openai_messages(
+    contents: list[types.Content], system_instruction: Optional[str] = None
+) -> list[ChatCompletionMessageParam]:
+    """Convert google.genai Content list to OpenAI messages format."""
+    messages: list[ChatCompletionMessageParam] = []
+
+    # Add system message if provided
+    if system_instruction:
+        system_message: ChatCompletionSystemMessageParam = {"role": "system", "content": system_instruction}
+        messages.append(system_message)
+
+    # Track tool calls to ensure proper flow
+    pending_tool_calls = set()
+
+    for content in contents:
+        role = _convert_role_to_openai(content.role)
+
+        # Separate different types of parts
+        text_parts = []
+        function_calls = []
+        function_responses = []
+        image_parts = []
+
+        for part in content.parts or []:
+            if part.text:
+                text_parts.append(part.text)
+            elif part.function_call:
+                function_calls.append(part)
+            elif part.function_response:
+                function_responses.append(part)
+            elif part.inline_data and part.inline_data.mime_type and part.inline_data.mime_type.startswith("image"):
+                if part.inline_data.data:
+                    image_data = base64.b64encode(part.inline_data.data).decode()
+                    image_part: ChatCompletionContentPartImageParam = {
+                        "type": "image_url",
+                        "image_url": {"url": f"data:{part.inline_data.mime_type};base64,{image_data}"},
+                    }
+                    image_parts.append(image_part)
+
+        # Handle function responses first (they should be tool messages)
+        for func_response in function_responses:
+            tool_call_id = func_response.function_response.id or "call_1"
+            if tool_call_id in pending_tool_calls:
+                tool_message: ChatCompletionToolMessageParam = {
+                    "role": "tool",
+                    "tool_call_id": tool_call_id,
+                    "content": str(func_response.function_response.response.get("result", ""))
+                    if func_response.function_response.response
+                    else "",
+                }
+                messages.append(tool_message)
+                pending_tool_calls.discard(tool_call_id)
+
+        # Handle function calls (assistant messages with tool_calls)
+        if function_calls:
+            tool_calls = []
+            for func_call in function_calls:
+                tool_call_function: ToolCallFunction = {
+                    "name": func_call.function_call.name or "",
+                    "arguments": str(func_call.function_call.args) if func_call.function_call.args else "{}",
+                }
+                tool_call_id = func_call.function_call.id or "call_1"
+                tool_call: ChatCompletionMessageToolCallParam = {
+                    "id": tool_call_id,
+                    "type": "function",
+                    "function": tool_call_function,
+                }
+                tool_calls.append(tool_call)
+                pending_tool_calls.add(tool_call_id)
+
+            # Create assistant message with tool calls
+            text_content = "\n".join(text_parts) if text_parts else None
+            assistant_message: ChatCompletionAssistantMessageParam = {
+                "role": "assistant",
+                "content": text_content,
+                "tool_calls": tool_calls,
+            }
+            messages.append(assistant_message)
+
+        # Handle regular text/image messages (only if no function calls)
+        elif text_parts or image_parts:
+            if role == "user":
+                if image_parts and text_parts:
+                    # Multi-modal content
+                    text_part: ChatCompletionContentPartTextParam = {"type": "text", "text": "\n".join(text_parts)}
+                    content_parts = [text_part] + image_parts
+                    user_message: ChatCompletionUserMessageParam = {"role": "user", "content": content_parts}
+                elif image_parts:
+                    # Image only
+                    user_message: ChatCompletionUserMessageParam = {"role": "user", "content": image_parts}
+                else:
+                    # Text only
+                    user_message: ChatCompletionUserMessageParam = {"role": "user", "content": "\n".join(text_parts)}
+                messages.append(user_message)
+            elif role == "assistant":
+                # Assistant messages with text (no tool calls)
+                assistant_message: ChatCompletionAssistantMessageParam = {
+                    "role": "assistant",
+                    "content": "\n".join(text_parts),
+                }
+                messages.append(assistant_message)
+
+    return messages
+
+
+def _update_type_string(value_dict: dict[str, Any]):
+    """Updates 'type' field to expected JSON schema format."""
+    if "type" in value_dict:
+        value_dict["type"] = value_dict["type"].lower()
+
+    if "items" in value_dict:
+        # 'type' field could exist for items as well, this would be the case if
+        # items represent primitive types.
+        _update_type_string(value_dict["items"])
+
+        if "properties" in value_dict["items"]:
+            # There could be properties as well on the items, especially if the items
+            # are complex object themselves. We recursively traverse each individual
+            # property as well and fix the "type" value.
+            for _, value in value_dict["items"]["properties"].items():
+                _update_type_string(value)
+
+    if "properties" in value_dict:
+        # Handle nested properties
+        for _, value in value_dict["properties"].items():
+            _update_type_string(value)
+
+
+def _convert_tools_to_openai(tools: list[types.Tool]) -> list[ChatCompletionToolParam]:
+    """Convert google.genai Tools to OpenAI tools format."""
+    openai_tools: list[ChatCompletionToolParam] = []
+
+    for tool in tools:
+        if tool.function_declarations:
+            for func_decl in tool.function_declarations:
+                # Build function definition
+                function_def: FunctionDefinition = {
+                    "name": func_decl.name or "",
+                    "description": func_decl.description or "",
+                }
+
+                # Always include parameters field, even if empty
+                properties = {}
+                required = []
+
+                if func_decl.parameters:
+                    if func_decl.parameters.properties:
+                        for prop_name, prop_schema in func_decl.parameters.properties.items():
+                            value_dict = prop_schema.model_dump(exclude_none=True)
+                            _update_type_string(value_dict)
+                            properties[prop_name] = value_dict
+
+                    if func_decl.parameters.required:
+                        required = func_decl.parameters.required
+
+                function_def["parameters"] = {"type": "object", "properties": properties, "required": required}
+
+                # Create the tool param
+                openai_tool: ChatCompletionToolParam = {"type": "function", "function": function_def}
+                openai_tools.append(openai_tool)
+
+    return openai_tools
+
+
+def _convert_openai_response_to_llm_response(response: ChatCompletion) -> LlmResponse:
+    """Convert OpenAI response to LlmResponse."""
+    choice = response.choices[0]
+    message = choice.message
+
+    parts = []
+
+    # Handle text content
+    if message.content:
+        parts.append(types.Part.from_text(text=message.content))
+
+    # Handle function calls
+    if hasattr(message, "tool_calls") and message.tool_calls:
+        for tool_call in message.tool_calls:
+            if tool_call.type == "function":
+                try:
+                    args = json.loads(tool_call.function.arguments) if tool_call.function.arguments else {}
+                except json.JSONDecodeError:
+                    args = {}
+
+                part = types.Part.from_function_call(name=tool_call.function.name, args=args)
+                if part.function_call:
+                    part.function_call.id = tool_call.id
+                parts.append(part)
+
+    content = types.Content(role="model", parts=parts)
+
+    # Handle usage metadata
+    usage_metadata = None
+    if hasattr(response, "usage") and response.usage:
+        usage_metadata = types.GenerateContentResponseUsageMetadata(
+            prompt_token_count=response.usage.prompt_tokens,
+            candidates_token_count=response.usage.completion_tokens,
+            total_token_count=response.usage.total_tokens,
+        )
+
+    # Handle finish reason
+    finish_reason = types.FinishReason.STOP
+    if choice.finish_reason == "length":
+        finish_reason = types.FinishReason.MAX_TOKENS
+    elif choice.finish_reason == "content_filter":
+        finish_reason = types.FinishReason.SAFETY
+
+    return LlmResponse(content=content, usage_metadata=usage_metadata, finish_reason=finish_reason)
+
+
+class BaseOpenAI(BaseLlm):
+    """Base class for OpenAI-compatible models."""
+
+    model: str
+    base_url: Optional[str] = None
+    api_key: Optional[str] = Field(default=None, exclude=True)
+    max_tokens: Optional[int] = None
+    temperature: Optional[float] = None
+
+    @classmethod
+    def supported_models(cls) -> list[str]:
+        """Returns a list of supported models in regex for LlmRegistry."""
+        return [r"gpt-.*", r"o1-.*"]
+
+    @cached_property
+    def _client(self) -> AsyncOpenAI:
+        """Get the OpenAI client."""
+        kwargs = {}
+        if self.base_url:
+            kwargs["base_url"] = self.base_url
+        if self.api_key:
+            kwargs["api_key"] = self.api_key
+
+        return AsyncOpenAI(**kwargs)
+
+    async def generate_content_async(
+        self, llm_request: LlmRequest, stream: bool = False
+    ) -> AsyncGenerator[LlmResponse, None]:
+        """Generate content using OpenAI API."""
+
+        # Convert messages
+        system_instruction = None
+        if llm_request.config and llm_request.config.system_instruction:
+            if isinstance(llm_request.config.system_instruction, str):
+                system_instruction = llm_request.config.system_instruction
+            elif hasattr(llm_request.config.system_instruction, "parts"):
+                # Handle Content type system instruction
+                text_parts = []
+                parts = getattr(llm_request.config.system_instruction, "parts", [])
+                if parts:
+                    for part in parts:
+                        if hasattr(part, "text") and part.text:
+                            text_parts.append(part.text)
+                    system_instruction = "\n".join(text_parts)
+
+        messages = _convert_content_to_openai_messages(llm_request.contents, system_instruction)
+
+        # Prepare request parameters
+        kwargs = {
+            "model": llm_request.model or self.model,
+            "messages": messages,
+        }
+
+        if self.max_tokens:
+            kwargs["max_tokens"] = self.max_tokens
+        if self.temperature is not None:
+            kwargs["temperature"] = self.temperature
+
+        # Handle tools
+        if llm_request.config and llm_request.config.tools:
+            # Filter to only google.genai.types.Tool objects
+            genai_tools = []
+            for tool in llm_request.config.tools:
+                if hasattr(tool, "function_declarations"):
+                    genai_tools.append(tool)
+
+            if genai_tools:
+                openai_tools = _convert_tools_to_openai(genai_tools)
+                if openai_tools:
+                    kwargs["tools"] = openai_tools
+                    kwargs["tool_choice"] = "auto"
+
+        try:
+            if stream:
+                # Handle streaming
+                async for chunk in await self._client.chat.completions.create(stream=True, **kwargs):
+                    if chunk.choices and chunk.choices[0].delta:
+                        delta = chunk.choices[0].delta
+                        if delta.content:
+                            content = types.Content(role="model", parts=[types.Part.from_text(text=delta.content)])
+                            yield LlmResponse(
+                                content=content, partial=True, turn_complete=chunk.choices[0].finish_reason is not None
+                            )
+            else:
+                # Handle non-streaming
+                response = await self._client.chat.completions.create(stream=False, **kwargs)
+                yield _convert_openai_response_to_llm_response(response)
+
+        except Exception as e:
+            yield LlmResponse(error_code="API_ERROR", error_message=str(e))
+
+
+class OpenAI(BaseOpenAI):
+    """OpenAI model implementation."""
+
+    type: Literal["openai"]
+
+    @cached_property
+    def _client(self) -> AsyncOpenAI:
+        """Get the OpenAI client."""
+        kwargs = {}
+        if self.base_url:
+            kwargs["base_url"] = self.base_url
+        if self.api_key:
+            kwargs["api_key"] = self.api_key
+        elif "OPENAI_API_KEY" in os.environ:
+            kwargs["api_key"] = os.environ["OPENAI_API_KEY"]
+
+        return AsyncOpenAI(**kwargs)
+
+
+class AzureOpenAI(BaseOpenAI):
+    """Azure OpenAI model implementation."""
+
+    type: Literal["azure_openai"]
+    api_version: Optional[str] = None
+    azure_endpoint: Optional[str] = None
+    azure_deployment: Optional[str] = None
+
+    @cached_property
+    def _client(self) -> AsyncAzureOpenAI:
+        """Get the Azure OpenAI client."""
+        api_version = self.api_version or os.environ.get("AZURE_OPENAI_API_VERSION", "2024-02-15-preview")
+        azure_endpoint = self.azure_endpoint or os.environ.get("AZURE_OPENAI_ENDPOINT")
+        api_key = self.api_key or os.environ.get("AZURE_OPENAI_API_KEY")
+
+        if not azure_endpoint:
+            raise ValueError(
+                "Azure endpoint must be provided either via azure_endpoint parameter or AZURE_OPENAI_ENDPOINT environment variable"
+            )
+
+        if not api_key:
+            raise ValueError(
+                "API key must be provided either via api_key parameter or AZURE_OPENAI_API_KEY environment variable"
+            )
+
+        return AsyncAzureOpenAI(api_version=api_version, azure_endpoint=azure_endpoint, api_key=api_key)