google adk using model from AI platform via endpoint

Sometimes we would like to make use of a deployed model in AI platform and this is what post intent to achieve.

We wrap our class in BaseLLM and then implement the relevant code in this case override the based method generate_content_async. And plug it into our root agent and use runner to execute this code.

This make it really easy to use.




import asyncio
from typing import AsyncGenerator
from google.cloud import aiplatform
from google.adk.agents import Agent
from google.adk.models import BaseLlm, LlmRequest, LlmResponse
from google.genai import types
from google.adk.runners import Runner
from google.adk.sessions import InMemorySessionService
from google.genai import types

class VertexEndpointLlm(BaseLlm):
    project: str
    location: str
    endpoint_id: str

    async def generate_content_async(
        self, llm_request: LlmRequest, stream: bool = False
    ) -> AsyncGenerator[LlmResponse, None]:
        # Flatten the ADK conversation into a single prompt
        prompt = "\n".join(
            p.text for c in llm_request.contents for p in (c.parts or []) if p.text
        )

        def _predict():
            aiplatform.init(project=self.project, location=self.location)
            ep = aiplatform.Endpoint(
                f"projects/{self.project}/locations/{self.location}/endpoints/{self.endpoint_id}"
            )
            return ep.predict(instances=[{"prompt": prompt}], use_dedicated_endpoint=True)

        resp = await asyncio.to_thread(_predict)   # SDK call is blocking
        text = str(resp.predictions[0])

        yield LlmResponse(content=types.Content(role="model", parts=[types.Part(text=text)]))

root_agent = Agent(
    name="endpoint_llm_agent",
    model=VertexEndpointLlm(
        model="custom-vertex-endpoint",   # label only
        project="YOUR-PROJECT-NUMBER",
        location="asia-southeast1",
        endpoint_id="mg-endpoint-DEPLOYED-TO-AI-PLATFORM",
    ),
    instruction="You are a helpful assistant.",
)

async def main():
    session_service = InMemorySessionService()
    runner = Runner(agent=root_agent, app_name="my_app", session_service=session_service)
    await session_service.create_session(app_name="my_app", user_id="u1", session_id="s1")

    message = types.Content(role="user", parts=[types.Part(text="Hello, summarise this: ...")])

    async for event in runner.run_async(user_id="u1", session_id="s1", new_message=message):
        if event.is_final_response() and event.content:
            print(event.content.parts[0].text)

asyncio.run(main())



Comments

Post a Comment

Popular posts from this blog

NodeJS: Error: spawn EINVAL in window for node version 20.20 and 18.20

Encounter this issue when running autorest (that uses nodejs).  I am using node version:  v18.20.2 After i revert to use node v18.12.0, then i was able to get my app to run without this error. Since i am working with Azure devops yaml, i was able to use the following to resolve my issue.  - task : NodeTool@0   inputs :     versionSource : 'spec'     versionSpec : '18.18.0'
Read more

llama cpp running it in google colab

Image
There are 2 options to do this in Google  colab.  Option 1 (Easiest)  With this option, we can just install the relevant package and then run the model.  First we need to install the required packages using the following command:-  ! pip install -U llama-cpp-python Next,  we will try to get the model using the following command # !pip install llama-cpp-python from llama_cpp import Llama llm = Llama.from_pretrained(   repo_id= "unsloth/Qwen3-0.6B-GGUF" ,   filename= "Qwen3-0.6B-IQ4_NL.gguf" , ) We can see the model being downloaded:- And then we will run the following command to test this model llm.create_chat_completion(   messages = [     {       "role" : "user" ,       "content" : "What is the capital of France?"     }   ] ) And the output will look something like this:-  Option 2 llama.cpp is a powerful inference engine and if you wanted to get it running in google colab, you co...
Read more