Improve docs with examples on custom agentchat agents with model clients and serialization

[victordibia]

This PR improves documentation on custom agents

• Shows example on how to create a custom agent that directly uses a model client. In this case an example of a GeminiAssistantAgent that directly uses the Gemini SDK model client.
• Shows that that CustomAgent can be easily added to any agentchat team
• Shows how the same CustomAgent can be made declarative by inheriting the Component interface and implementing the required methods.

import os
from typing import AsyncGenerator, Sequence

from autogen_agentchat.agents import BaseChatAgent
from autogen_agentchat.base import Response
from autogen_agentchat.messages import AgentEvent, ChatMessage
from autogen_core import CancellationToken, Component
from pydantic import BaseModel


class GeminiAssistantAgentConfig(BaseModel):
    name: str
    description: str = "An agent that provides assistance with ability to use tools."
    model: str = "gemini-1.5-flash-002"
    system_message: str | None = None


class GeminiAssistant(BaseChatAgent, Component[GeminiAssistantAgentConfig]):
    component_config_schema = GeminiAssistantAgentConfig
    # component_provider_override = "mypackage.agents.GeminiAssistant"

    def __init__(
        self,
        name: str,
        description: str = "An agent that provides assistance with ability to use tools.",
        model: str = "gemini-1.5-flash-002",
        api_key: str = os.environ["GEMINI_API_KEY"],
        system_message: str = "You are a helpful assistant that can respond to messages. Reply with TERMINATE when the task has been completed.",
    ):
        super().__init__(name=name, description=description)
        self._model_context = UnboundedChatCompletionContext()
        self._model_client = genai.Client(api_key=api_key)
        self._system_message = system_message
        self._model = model

    @property
    def produced_message_types(self) -> Sequence[type[ChatMessage]]:
        return (TextMessage,)

    async def on_messages(self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken) -> Response:
        async for message in self.on_messages_stream(messages, cancellation_token):
            if isinstance(message, Response):
                return message
        raise AssertionError("The stream should have returned the final result.")

    async def on_messages_stream(
        self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken
    ) -> AsyncGenerator[AgentEvent | ChatMessage | Response, None]:
        # Add messages to the model context
        for msg in messages:
            await self._model_context.add_message(UserMessage(content=msg.content, source=msg.source))

        # Get conversation history
        history = [msg.source + ": " + msg.content + "\n" for msg in await self._model_context.get_messages()]

        # Generate response using Gemini
        response = self._model_client.models.generate_content(
            model=self._model,
            contents=f"History: {history}\nGiven the history, please provide a response",
            config=types.GenerateContentConfig(
                system_instruction=self._system_message,
                temperature=0.3,
            ),
        )

        # Create usage metadata
        usage = RequestUsage(
            prompt_tokens=response.usage_metadata.prompt_token_count,
            completion_tokens=response.usage_metadata.candidates_token_count,
        )

        # Add response to model context
        await self._model_context.add_message(AssistantMessage(content=response.text, source=self.name))

        # Yield the final response
        yield Response(
            chat_message=TextMessage(content=response.text, source=self.name, models_usage=usage),
            inner_messages=[],
        )

    async def on_reset(self, cancellation_token: CancellationToken) -> None:
        """Reset the assistant by clearing the model context."""
        await self._model_context.clear()

    @classmethod
    def _from_config(cls, config: GeminiAssistantAgentConfig) -> "GeminiAssistant":
        return cls(
            name=config.name, description=config.description, model=config.model, system_message=config.system_message
        )

    def _to_config(self) -> GeminiAssistantAgentConfig:
        return GeminiAssistantAgentConfig(
            name=self.name,
            description=self.description,
            model=self._model,
            system_message=self._system_message,
        )

gemini_assistant = GeminiAssistant("gemini_assistant")
config = gemini_assistant.dump_component()
print(config.model_dump_json(indent=2))
loaded_agent = GeminiAssistant.load_component(config)
print(loaded_agent)

team = RoundRobinGroupChat([primary_agent, gemini_critic_agent], termination_condition=text_termination)

Why are these changes needed?

Related issue number

Closes #5450

Checks

• I've included any doc changes needed for https://microsoft.github.io/autogen/. See https://microsoft.github.io/autogen/docs/Contribute#documentation to build and test documentation locally.
• I've added tests (if relevant) corresponding to the changes introduced in this PR.
• I've made sure all auto checks have passed.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 78.33%. Comparing base (340a8e8) to head (eb36cda).

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #5468      +/-   ##
==========================================
+ Coverage   78.32%   78.33%   +0.01%     
==========================================
  Files         165      165              
  Lines        9800     9800              
==========================================
+ Hits         7676     7677       +1     
+ Misses       2124     2123       -1     

Flag	Coverage Δ
unittests	78.33% <ø> (+0.01%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[ekzhu]

we can move the location of custom-agents.ipynb to the parent folder. Then add a new redirect entry in the conf.py: https://github.com/microsoft/autogen/blob/main/python/packages/autogen-core/docs/src/conf.py#L175-L176

[ekzhu]

I think generally it is a good idea to pass the gemini client through the constructor. This is the convention we use for other agents, also it allows the client to be shared.

[ekzhu]

Also we need to update the uv.lock file as now the notebook depends on the google package.

[ekzhu]

I think generally it is a good idea to pass the gemini client through the constructor. This is the convention we use for other agents, also it allows the client to be shared.

Oh I see, you did this because the gemini client isn't serializable... I guess it's fine then.

[ekzhu]

Returnning within the generator can cause the generator itself not properly closed -- it has async tasks associated with it. Let the generator context close first then return the result.

[ekzhu]

I think Gemini supports chat messages right? https://ai.google.dev/gemini-api/docs/text-generation?lang=python#chat

For this example we don't necessarly need to use model context here. If Gemini can manage the context -- it's 1M token, we can just simplify and use the Gemini's chat feature directly.

[ekzhu]

We can write in Next Step to encourage reader to try creating a custom client with function calling capability following the Gemini's API doc: https://ai.google.dev/gemini-api/docs/function-calling