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

kagent/adk/__init__.py

@@ -1,7 +1,7 @@
 import importlib.metadata
 
 from ._a2a import KAgentApp
-from .models import AgentConfig
+from .types import AgentConfig
 
 __version__ = importlib.metadata.version("kagent_adk")
 

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)

kagent/adk/{models.py → types.py}

@@ -12,6 +12,9 @@ from google.adk.tools.agent_tool import AgentTool
 from google.adk.tools.mcp_tool import MCPToolset, SseConnectionParams, StreamableHTTPConnectionParams
 from pydantic import BaseModel, Field
 
+from .models import AzureOpenAI as OpenAIAzure
+from .models import OpenAI as OpenAINative
+
 logger = logging.getLogger(__name__)
 
 
@@ -97,7 +100,7 @@ class AgentConfig(BaseModel):
                 mcp_toolsets.append(AgentTool(agent=remote_agent, skip_summarization=True))
 
         if self.model.type == "openai":
-            model = LiteLlm(model=f"openai/{self.model.model}", base_url=self.model.base_url)
+            model = OpenAINative(model=self.model.model, base_url=self.model.base_url, type="openai")
         elif self.model.type == "anthropic":
             model = LiteLlm(model=f"anthropic/{self.model.model}", base_url=self.model.base_url)
         elif self.model.type == "gemini_vertex_ai":
@@ -107,7 +110,7 @@ class AgentConfig(BaseModel):
         elif self.model.type == "ollama":
             model = LiteLlm(model=f"ollama_chat/{self.model.model}")
         elif self.model.type == "azure_openai":
-            model = LiteLlm(model=f"azure/{self.model.model}")
+            model = OpenAIAzure(model=self.model.model, type="azure_openai")
         elif self.model.type == "gemini":
             model = self.model.model
         else:

{kagent_adk-0.6.8.dist-info → kagent_adk-0.6.9.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kagent-adk
-Version: 0.6.8
+Version: 0.6.9
 Summary: kagent-adk is an sdk for integrating adk agents with kagent
 Requires-Python: >=3.12.11
 Requires-Dist: a2a-sdk>=0.3.1

{kagent_adk-0.6.8.dist-info → kagent_adk-0.6.9.dist-info}/RECORD

@@ -1,15 +1,17 @@
-kagent/adk/__init__.py,sha256=oTnqo17F9mqzSfxs0-b_AajVSxc17KCuwHleuxdUzH4,183
+kagent/adk/__init__.py,sha256=Cam9hwhl6Z1tZ3WJIB1KATX24MwRMpiyBLYF64DTqQI,182
 kagent/adk/_a2a.py,sha256=HvduKJlfuUIKzYuPs_U_9A-xuhcMdQutZTzdLNaGORc,4406
 kagent/adk/_agent_executor.py,sha256=rxNVdrcS4L5Kwe8oF7mLUnD1KOYhe6vXK-si0O-KQ9o,10730
 kagent/adk/_session_service.py,sha256=A47gsfDVp8jITzeW987AHTJLEhcU_mU3ik_SFptFGIc,5815
 kagent/adk/_token.py,sha256=OL46m7U5vUTby1WWjVB7Jqzig4TWddzoAmLVLlfSdAg,2515
 kagent/adk/cli.py,sha256=Sdw9FXnUmeHfgJhdRDq1zpoNMf6GM7x35ZDz_OkQgJg,2963
-kagent/adk/models.py,sha256=C_h8dxh_PKEXRM3kSMg5YkIxA2kbagGPfKMbpnLTZj8,4147
+kagent/adk/types.py,sha256=v7VbO0tAPvpjdEgWJroEq64jSCz6D3Q0yPdRM0a7MjM,4260
 kagent/adk/converters/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 kagent/adk/converters/event_converter.py,sha256=gPDxJhPtoIp2CAm0VjSPTH-siuz3IR23CV25AiN_1B8,10392
 kagent/adk/converters/part_converter.py,sha256=Q9Kbteit8XdL_9Tb8bAtxRxOZQei8Wabszwsl4YMe2c,7507
 kagent/adk/converters/request_converter.py,sha256=iTmTmhlnyRfuYyFi4WmpTSXPz22xjjotbe750j-CvYA,1072
-kagent_adk-0.6.8.dist-info/METADATA,sha256=SY_f2O3Y7EcWfjH33sdFLsSHpKladu5QKMT1PLJEgw8,943
-kagent_adk-0.6.8.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-kagent_adk-0.6.8.dist-info/entry_points.txt,sha256=a1Q2Inc9L0dvXWEkwnCdf9cfXdpX5Dl2Q6DhNWNjhxw,50
-kagent_adk-0.6.8.dist-info/RECORD,,
+kagent/adk/models/__init__.py,sha256=mqD0JhS9kT1rMpFNLq5-qnjstpp6lzT9xADaOfjrUKY,78
+kagent/adk/models/_openai.py,sha256=Xdu0AteM4nFbZaCMAcSC_nj1iJRgpWBVXFzD0aPELGA,15685
+kagent_adk-0.6.9.dist-info/METADATA,sha256=JtlHtn8MXVr9XeCg8QX3C0zob7ibyNhf1N6ygfQTLlE,943
+kagent_adk-0.6.9.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+kagent_adk-0.6.9.dist-info/entry_points.txt,sha256=a1Q2Inc9L0dvXWEkwnCdf9cfXdpX5Dl2Q6DhNWNjhxw,50
+kagent_adk-0.6.9.dist-info/RECORD,,