99 lines
3.4 KiB
Python
99 lines
3.4 KiB
Python
import os
|
|
import asyncio
|
|
from typing import Optional
|
|
|
|
import agentops
|
|
from google.adk.agents import LlmAgent
|
|
from google.adk.tools import FunctionTool
|
|
from google.adk.runners import Runner
|
|
from google.adk.sessions import InMemorySessionService
|
|
from google.genai import types
|
|
|
|
|
|
# Attempt to import RunConfig/StreamingMode from likely ADK locations
|
|
RunConfig: Optional[object] = None
|
|
StreamingMode: Optional[object] = None
|
|
try:
|
|
from google.adk.runners import RunConfig as _RunConfig, StreamingMode as _StreamingMode # type: ignore
|
|
|
|
RunConfig = _RunConfig
|
|
StreamingMode = _StreamingMode
|
|
except Exception:
|
|
try:
|
|
from google.adk.types import RunConfig as _RunConfig2, StreamingMode as _StreamingMode2 # type: ignore
|
|
|
|
RunConfig = _RunConfig2
|
|
StreamingMode = _StreamingMode2
|
|
except Exception:
|
|
RunConfig = None
|
|
StreamingMode = None
|
|
|
|
|
|
# Initialize AgentOps (set AGENTOPS_API_KEY in your environment)
|
|
agentops.init(api_key=os.getenv("AGENTOPS_API_KEY"), trace_name="adk_sse_text_function_call")
|
|
|
|
APP_NAME = "adk_sse_text_function_call_app"
|
|
USER_ID = "user_sse_text_fc"
|
|
SESSION_ID = "session_sse_text_fc"
|
|
MODEL_NAME = "gemini-2.0-flash"
|
|
|
|
|
|
# Simple tool to trigger a function call
|
|
async def get_weather(location: str) -> str:
|
|
return f"Weather for {location}: sunny and 25°C."
|
|
|
|
|
|
weather_tool = FunctionTool(func=get_weather)
|
|
|
|
# Agent configured with the tool so the model can trigger a function call
|
|
agent = LlmAgent(
|
|
model=MODEL_NAME,
|
|
name="WeatherAgent",
|
|
description="Provides weather using a tool",
|
|
instruction=(
|
|
"You are a helpful assistant. When asked about weather, call the get_weather tool with the given location."
|
|
),
|
|
tools=[weather_tool],
|
|
output_key="weather_output",
|
|
)
|
|
|
|
# Session service and runner
|
|
session_service = InMemorySessionService()
|
|
runner = Runner(agent=agent, app_name=APP_NAME, session_service=session_service)
|
|
|
|
|
|
async def main():
|
|
# Ensure session exists
|
|
await session_service.create_session(app_name=APP_NAME, user_id=USER_ID, session_id=SESSION_ID)
|
|
|
|
# Create user message
|
|
user_message = types.Content(role="user", parts=[types.Part(text="What's the weather in Paris?")])
|
|
|
|
# Configure SSE streaming with TEXT modality, as reported by the user
|
|
run_config_kw = {}
|
|
if RunConfig is not None and StreamingMode is not None:
|
|
run_config_kw["run_config"] = RunConfig(streaming_mode=StreamingMode.SSE, response_modalities=["TEXT"]) # type: ignore
|
|
|
|
final_text = None
|
|
async for event in runner.run_async(
|
|
user_id=USER_ID, session_id=SESSION_ID, new_message=user_message, **run_config_kw
|
|
):
|
|
# Print out any parts safely; this will include function_call parts when they occur
|
|
if hasattr(event, "content") and event.content and getattr(event.content, "parts", None):
|
|
for part in event.content.parts:
|
|
text = getattr(part, "text", None)
|
|
func_call = getattr(part, "function_call", None)
|
|
if text:
|
|
print(f"Assistant: {text}")
|
|
final_text = text
|
|
elif func_call is not None:
|
|
name = getattr(func_call, "name", "<unknown>")
|
|
args = getattr(func_call, "args", {})
|
|
print(f"Function call: {name} args={args}")
|
|
|
|
print("Final text:", final_text)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main())
|