from pydantic import BaseModel
from typing import Literal
# The response format for the agent as a Pydantic base model.
class AgentResponse(BaseModel):
    thoughts: str
    response: Literal["happy", "sad", "neutral"]

model_client = OpenAIChatCompletionClient(
    model="gemini-2.0-flash-exp",
    base_url="https://generativelanguage.googleapis.com/v1beta/openai/",
    api_key="AIzaSyCnEtytHlZlF1Z0J7-Gx3scxNX8xaw-9Qc",
    model_info={
        "vision": True,
        "function_calling": True,
        "json_output": True,
        "family": "unknown",
    },
    response_format=AgentResponse,
)
agent = AssistantAgent(
    name="Agent",
    model_client=model_client,
    tools=tools,
    #system_message= "You are an assistant with some tools that can be used to answer some questions",
    system_message="Categorize the input as happy, sad, or neutral following the JSON format.",
    model_client_stream=True,
)
async for message in agent.on_messages_stream([TextMessage(content="I am happy.", source="user")], CancellationToken()):
    print(f"{message.}", end=" ")