97 | No proverbs yet. Ask the assistant to add some!
98 |
}
99 |
100 |
101 | );
102 | }
103 |
104 | // Simple sun icon for the weather card
105 | function SunIcon() {
106 | return (
107 |
111 | );
112 | }
113 |
114 | // Weather card component where the location and themeColor are based on what the agent
115 | // sets via tool calls.
116 | function WeatherCard({ location, themeColor }: { location?: string, themeColor: string }) {
117 | return (
118 |
122 |
123 |
124 |
125 |
{location}
126 |
Current Weather
127 |
128 |
129 |
130 |
131 |
132 |
70°
133 |
Clear skies
134 |
135 |
136 |
137 |
138 |
139 |
Humidity
140 |
45%
141 |
142 |
143 |
Wind
144 |
5 mph
145 |
146 |
147 |
Feels Like
148 |
72°
149 |
150 |
151 |
152 |
153 |
154 | );
155 | }
156 |
--------------------------------------------------------------------------------
/agent/agent.py:
--------------------------------------------------------------------------------
1 | """Shared State feature."""
2 |
3 | from __future__ import annotations
4 |
5 | from dotenv import load_dotenv
6 | load_dotenv()
7 | import json
8 | from enum import Enum
9 | from typing import Dict, List, Any, Optional
10 | from fastapi import FastAPI
11 | from ag_ui_adk import ADKAgent, add_adk_fastapi_endpoint
12 |
13 | # ADK imports
14 | from google.adk.agents import LlmAgent
15 | from google.adk.agents.callback_context import CallbackContext
16 | from google.adk.sessions import InMemorySessionService, Session
17 | from google.adk.runners import Runner
18 | from google.adk.events import Event, EventActions
19 | from google.adk.tools import FunctionTool, ToolContext
20 | from google.genai.types import Content, Part , FunctionDeclaration
21 | from google.adk.models import LlmResponse, LlmRequest
22 | from google.genai import types
23 |
24 | from pydantic import BaseModel, Field
25 | from typing import List, Optional
26 | from enum import Enum
27 |
28 |
29 | class ProverbsState(BaseModel):
30 | """List of the proverbs being written."""
31 | proverbs: list[str] = Field(
32 | default_factory=list,
33 | description='The list of already written proverbs',
34 | )
35 |
36 |
37 | def set_proverbs(
38 | tool_context: ToolContext,
39 | new_proverbs: list[str]
40 | ) -> Dict[str, str]:
41 | """
42 | Set the list of provers using the provided new list.
43 |
44 | Args:
45 | "new_proverbs": {
46 | "type": "array",
47 | "items": {"type": "string"},
48 | "description": "The new list of proverbs to maintain",
49 | }
50 |
51 | Returns:
52 | Dict indicating success status and message
53 | """
54 | try:
55 | # Put this into a state object just to confirm the shape
56 | new_state = { "proverbs": new_proverbs}
57 | tool_context.state["proverbs"] = new_state["proverbs"]
58 | return {"status": "success", "message": "Proverbs updated successfully"}
59 |
60 | except Exception as e:
61 | return {"status": "error", "message": f"Error updating proverbs: {str(e)}"}
62 |
63 |
64 |
65 | def get_weather(tool_context: ToolContext, location: str) -> Dict[str, str]:
66 | """Get the weather for a given location. Ensure location is fully spelled out."""
67 | return {"status": "success", "message": f"The weather in {location} is sunny."}
68 |
69 |
70 |
71 | def on_before_agent(callback_context: CallbackContext):
72 | """
73 | Initialize proverbs state if it doesn't exist.
74 | """
75 |
76 | if "proverbs" not in callback_context.state:
77 | # Initialize with default recipe
78 | default_proverbs = []
79 | callback_context.state["proverbs"] = default_proverbs
80 |
81 |
82 | return None
83 |
84 |
85 |
86 |
87 | # --- Define the Callback Function ---
88 | # modifying the agent's system prompt to incude the current state of the proverbs list
89 | def before_model_modifier(
90 | callback_context: CallbackContext, llm_request: LlmRequest
91 | ) -> Optional[LlmResponse]:
92 | """Inspects/modifies the LLM request or skips the call."""
93 | agent_name = callback_context.agent_name
94 | if agent_name == "ProverbsAgent":
95 | proverbs_json = "No proverbs yet"
96 | if "proverbs" in callback_context.state and callback_context.state["proverbs"] is not None:
97 | try:
98 | proverbs_json = json.dumps(callback_context.state["proverbs"], indent=2)
99 | except Exception as e:
100 | proverbs_json = f"Error serializing proverbs: {str(e)}"
101 | # --- Modification Example ---
102 | # Add a prefix to the system instruction
103 | original_instruction = llm_request.config.system_instruction or types.Content(role="system", parts=[])
104 | prefix = f"""You are a helpful assistant for maintaining a list of proverbs.
105 | This is the current state of the list of proverbs: {proverbs_json}
106 | When you modify the list of proverbs (wether to add, remove, or modify one or more proverbs), use the set_proverbs tool to update the list."""
107 | # Ensure system_instruction is Content and parts list exists
108 | if not isinstance(original_instruction, types.Content):
109 | # Handle case where it might be a string (though config expects Content)
110 | original_instruction = types.Content(role="system", parts=[types.Part(text=str(original_instruction))])
111 | if not original_instruction.parts:
112 | original_instruction.parts.append(types.Part(text="")) # Add an empty part if none exist
113 |
114 | # Modify the text of the first part
115 | modified_text = prefix + (original_instruction.parts[0].text or "")
116 | original_instruction.parts[0].text = modified_text
117 | llm_request.config.system_instruction = original_instruction
118 |
119 |
120 |
121 | return None
122 |
123 |
124 |
125 |
126 |
127 |
128 | # --- Define the Callback Function ---
129 | def simple_after_model_modifier(
130 | callback_context: CallbackContext, llm_response: LlmResponse
131 | ) -> Optional[LlmResponse]:
132 | """Stop the consecutive tool calling of the agent"""
133 | agent_name = callback_context.agent_name
134 | # --- Inspection ---
135 | if agent_name == "ProverbsAgent":
136 | original_text = ""
137 | if llm_response.content and llm_response.content.parts:
138 | # Assuming simple text response for this example
139 | if llm_response.content.role=='model' and llm_response.content.parts[0].text:
140 | original_text = llm_response.content.parts[0].text
141 | callback_context._invocation_context.end_invocation = True
142 |
143 | elif llm_response.error_message:
144 | return None
145 | else:
146 | return None # Nothing to modify
147 | return None
148 |
149 |
150 | proverbs_agent = LlmAgent(
151 | name="ProverbsAgent",
152 | model="gemini-2.5-flash",
153 | instruction=f"""
154 | When a user asks you to do anything regarding proverbs, you MUST use the set_proverbs tool.
155 |
156 | IMPORTANT RULES ABOUT PROVERBS AND THE SET_PROVERBS TOOL:
157 | 1. Always use the set_proverbs tool for any proverbs-related requests
158 | 2. Always pass the COMPLETE LIST of proverbs to the set_proverbs tool. If the list had 5 proverbs and you removed one, you must pass the complete list of 4 remaining proverbs.
159 | 3. You can use existing proverbs if one is relevant to the user's request, but you can also create new proverbs as required.
160 | 4. Be creative and helpful in generating complete, practical proverbs
161 | 5. After using the tool, provide a brief summary of what you create, removed, or changed 7.
162 |
163 | Examples of when to use the set_proverbs tool:
164 | - "Add a proverb about soap" → Use tool with an array containing the existing list of proverbs with the new proverb about soap at the end.
165 | - "Remove the first proverb" → Use tool with an array containing the all of the existing proverbs except the first one"
166 | - "Change any proverbs about cats to mention that they have 18 lives" → If no proverbs mention cats, do not use the tool. If one or more proverbs do mention cats, change them to mention cats having 18 lives, and use the tool with an array of all of the proverbs, including ones that were changed and ones that did not require changes.
167 |
168 | Do your best to ensure proverbs plausibly make sense.
169 |
170 |
171 | IMPORTANT RULES ABOUT WEATHER AND THE GET_WEATHER TOOL:
172 | 1. Only call the get_weather tool if the user asks you for the weather in a given location.
173 | 2. If the user does not specify a location, you can use the location "Everywhere ever in the whole wide world"
174 |
175 | Examples of when to use the get_weather tool:
176 | - "What's the weather today in Tokyo?" → Use the tool with the location "Tokyo"
177 | - "Whats the weather right now" → Use the location "Everywhere ever in the whole wide world"
178 | - Is it raining in London? → Use the tool with the location "London"
179 | """,
180 | tools=[set_proverbs, get_weather],
181 | before_agent_callback=on_before_agent,
182 | before_model_callback=before_model_modifier,
183 | after_model_callback = simple_after_model_modifier
184 | )
185 |
186 | # Create ADK middleware agent instance
187 | adk_proverbs_agent = ADKAgent(
188 | adk_agent=proverbs_agent,
189 | app_name="proverbs_app",
190 | user_id="demo_user",
191 | session_timeout_seconds=3600,
192 | use_in_memory_services=True
193 | )
194 |
195 | # Create FastAPI app
196 | app = FastAPI(title="ADK Middleware Proverbs Agent")
197 |
198 | # Add the ADK endpoint
199 | add_adk_fastapi_endpoint(app, adk_proverbs_agent, path="/")
200 |
201 | if __name__ == "__main__":
202 | import os
203 | import uvicorn
204 |
205 | if not os.getenv("GOOGLE_API_KEY"):
206 | print("⚠️ Warning: GOOGLE_API_KEY environment variable not set!")
207 | print(" Set it with: export GOOGLE_API_KEY='your-key-here'")
208 | print(" Get a key from: https://makersuite.google.com/app/apikey")
209 | print()
210 |
211 | port = int(os.getenv("PORT", 8000))
212 | uvicorn.run(app, host="0.0.0.0", port=port)
213 |
--------------------------------------------------------------------------------