Implementing a custom LiteLlm for Google agent development kit that supports tool calls

import os
import json
import aiohttp
import logging
from typing import Any, Dict, AsyncGenerator, Optional
from google.adk.models.basellm import BaseLlm
from google.adk.models.llmrequest import LlmRequest
from google.adk.models.llmresponse import LlmResponse
from google.genai import types  # ADK uses this for structured parts
logger = logging.getLogger(name)
class MyCustomClient:
    """Minimal HTTP client for a custom LLM endpoint."""
    def init(self, baseurl: str, apikey: str):
        self.baseurl = baseurl.rstrip("/")
        self.apikey = apikey
        self.session: Optional[aiohttp.ClientSession] = None
    async def getsession(self) -> aiohttp.ClientSession:
        if self.session is None or self.session.closed:
            self.session = aiohttp.ClientSession()
        return self.session
    async def close(self):
        if self.session and not self.session.closed:
            await self.session.close()
    async def generate(self, payload: Dict[str, Any]) -> Dict[str, Any]:
        session = await self.getsession()
        headers = {"Authorization": f"Bearer {self.apikey}"}
        url = f"{self.baseurl}/v1/chat/completions"
        async with session.post(url, json=payload, headers=headers) as response:
            response.raiseforstatus()
            return await response.json()
class MyLiteLlm(BaseLlm):
    """Custom LLM implementation that supports tool calls."""
    def init(self, model: str, baseurl: str, apikey: str, kwargs):
        super().init(model=model, kwargs)
        self.client = MyCustomClient(baseurl, apikey)
        self.model = model
    async def generatecontentasync(
        self, llmrequest: LlmRequest, stream: bool = False
    ) -> AsyncGenerator[LlmResponse, None]:
        if not llmrequest.contents:
            raise ValueError("LlmRequest must contain contents")
        payload = {
            "model": self.model,
            "messages": [c.modeldump(excludenone=True) for c in llmrequest.contents],
        }
        # optional: include tool declarations if present
        if llmrequest.tools:
            payload["tools"] = [t.modeldump(excludenone=True) for t in llmrequest.tools]
        apiresponse = await self.client.generate(payload)
        # 🔑 Convert response into ADK's LlmResponse, preserving tool calls
        parts = []
        for choice in apiresponse.get("choices", []):
            msg = choice.get("message", {})
            if "content" in msg and msg["content"]:
                parts.append(types.Part.fromtext(text=msg["content"]))
            if "toolcalls" in msg:
                for tc in msg["toolcalls"]:
                    parts.append(
                        types.Part(
                            functioncall=types.FunctionCall(
                                name=tc.get("function", {}).get("name", ""),
                                args=tc.get("function", {}).get("arguments", {}),
                            )
                        )
                    )
        llmresponse = LlmResponse(
            content=types.Content(
                role=apiresponse.get("role", "model"),
                parts=parts,
            ),
            partial=False,
        )
        yield llmresponse
    async def aenter(self):
        return self
    async def aexit(self, exctype, excval, exctb):
        await self.client.close()