fix: ADK 1.x Runner pattern + gemini-3.5-flash for A+ reasoning
This commit is contained in:
parent
497509370e
commit
1cba5166db
|
|
@ -4,7 +4,8 @@ agent.py — OSVauco root agent using ADK 1.x with A / A+ mode routing.
|
||||||
|
|
||||||
Modes:
|
Modes:
|
||||||
A (standard) — gemini-2.0-flash, $1/task hard stop
|
A (standard) — gemini-2.0-flash, $1/task hard stop
|
||||||
A+ (audit+) — gemini-2.5-pro orchestrator, $3/task hard stop,
|
A+ (audit+) — gemini-2.5-pro orchestrator + subagents,
|
||||||
|
gemini-3.5-flash reasoning, $3/task hard stop,
|
||||||
multi-agent pipeline enabled
|
multi-agent pipeline enabled
|
||||||
|
|
||||||
Authorized users for A+: opax, admin
|
Authorized users for A+: opax, admin
|
||||||
|
|
@ -13,11 +14,15 @@ Requires: google-adk >= 1.0.0,<2.0.0
|
||||||
google-cloud-aiplatform >= 1.112.0
|
google-cloud-aiplatform >= 1.112.0
|
||||||
"""
|
"""
|
||||||
|
|
||||||
|
import asyncio
|
||||||
import os
|
import os
|
||||||
import logging
|
import logging
|
||||||
from typing import Literal
|
from typing import Literal
|
||||||
|
|
||||||
from google.adk.agents import Agent
|
from google.adk.agents import Agent
|
||||||
|
from google.adk.runners import Runner
|
||||||
|
from google.adk.sessions.in_memory_session_service import InMemorySessionService
|
||||||
|
from google.genai import types
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
@ -37,7 +42,7 @@ REASONING_MODEL = os.environ.get("REASONING_MODEL", "gemini-2.0-flash")
|
||||||
# A+ (audit+) models
|
# A+ (audit+) models
|
||||||
HEAVY_ORCHESTRATOR = os.environ.get("HEAVY_ORCHESTRATOR_MODEL", "gemini-2.5-pro")
|
HEAVY_ORCHESTRATOR = os.environ.get("HEAVY_ORCHESTRATOR_MODEL", "gemini-2.5-pro")
|
||||||
HEAVY_SUBAGENT = os.environ.get("HEAVY_SUBAGENT_MODEL", "gemini-2.5-pro")
|
HEAVY_SUBAGENT = os.environ.get("HEAVY_SUBAGENT_MODEL", "gemini-2.5-pro")
|
||||||
HEAVY_REASONING = os.environ.get("HEAVY_REASONING_MODEL", "gemini-2.5-flash")
|
HEAVY_REASONING = os.environ.get("HEAVY_REASONING_MODEL", "gemini-3.5-flash") # NOT 2.5-flash
|
||||||
|
|
||||||
# Budget hard stops (USD per task)
|
# Budget hard stops (USD per task)
|
||||||
BUDGET_A = float(os.environ.get("BUDGET_A_USD_PER_TASK", "1.0"))
|
BUDGET_A = float(os.environ.get("BUDGET_A_USD_PER_TASK", "1.0"))
|
||||||
|
|
@ -46,6 +51,7 @@ BUDGET_APLUS = float(os.environ.get("HEAVY_MODE_BUDGET_USD_PER_DAY", "3.
|
||||||
HEAVY_MODE_ALLOWED_USERS = ["opax", "admin"]
|
HEAVY_MODE_ALLOWED_USERS = ["opax", "admin"]
|
||||||
|
|
||||||
Mode = Literal["A", "A+"]
|
Mode = Literal["A", "A+"]
|
||||||
|
APP_NAME = "opax"
|
||||||
|
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
# ---------------------------------------------------------------------------
|
||||||
|
|
@ -138,9 +144,48 @@ root_agent = build_agent(mode="A")
|
||||||
|
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
# ---------------------------------------------------------------------------
|
||||||
# run() — programmatic entry point for Cloud Run /run endpoint
|
# run() — async core, sync wrapper for Cloud Run /run endpoint
|
||||||
# ---------------------------------------------------------------------------
|
# ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
async def _run_async(
|
||||||
|
message: str,
|
||||||
|
user_id: str,
|
||||||
|
session_id: str,
|
||||||
|
mode: Mode,
|
||||||
|
) -> str:
|
||||||
|
"""Async implementation using ADK 1.x Runner pattern."""
|
||||||
|
agent = build_agent(mode=mode)
|
||||||
|
session_service = InMemorySessionService()
|
||||||
|
|
||||||
|
session = await session_service.create_session(
|
||||||
|
app_name=APP_NAME,
|
||||||
|
user_id=user_id,
|
||||||
|
session_id=session_id,
|
||||||
|
)
|
||||||
|
|
||||||
|
runner = Runner(
|
||||||
|
agent=agent,
|
||||||
|
app_name=APP_NAME,
|
||||||
|
session_service=session_service,
|
||||||
|
)
|
||||||
|
|
||||||
|
new_message = types.Content(
|
||||||
|
role="user",
|
||||||
|
parts=[types.Part(text=message)],
|
||||||
|
)
|
||||||
|
|
||||||
|
final_text = ""
|
||||||
|
async for event in runner.run_async(
|
||||||
|
user_id=user_id,
|
||||||
|
session_id=session.id,
|
||||||
|
new_message=new_message,
|
||||||
|
):
|
||||||
|
if event.is_final_response() and event.content and event.content.parts:
|
||||||
|
final_text = event.content.parts[0].text or ""
|
||||||
|
|
||||||
|
return final_text
|
||||||
|
|
||||||
|
|
||||||
def run(
|
def run(
|
||||||
message: str,
|
message: str,
|
||||||
user_id: str = "opax",
|
user_id: str = "opax",
|
||||||
|
|
@ -160,7 +205,7 @@ def run(
|
||||||
Agent response as a string.
|
Agent response as a string.
|
||||||
|
|
||||||
Raises:
|
Raises:
|
||||||
PermissionError: If user_id is not authorized for the requested mode.
|
PermissionError: If user_id is not authorized for A+ mode.
|
||||||
ValueError: If mode is not 'A' or 'A+'.
|
ValueError: If mode is not 'A' or 'A+'.
|
||||||
"""
|
"""
|
||||||
if mode not in ("A", "A+"):
|
if mode not in ("A", "A+"):
|
||||||
|
|
@ -168,13 +213,12 @@ def run(
|
||||||
|
|
||||||
authorize_mode(user_id, mode)
|
authorize_mode(user_id, mode)
|
||||||
|
|
||||||
agent = build_agent(mode=mode)
|
return asyncio.run(_run_async(
|
||||||
response = agent.run(
|
|
||||||
message=message,
|
message=message,
|
||||||
session_id=session_id,
|
|
||||||
user_id=user_id,
|
user_id=user_id,
|
||||||
)
|
session_id=session_id,
|
||||||
return response.text if hasattr(response, "text") else str(response)
|
mode=mode,
|
||||||
|
))
|
||||||
|
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
# ---------------------------------------------------------------------------
|
||||||
|
|
@ -183,6 +227,7 @@ def run(
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
import sys
|
import sys
|
||||||
|
logging.basicConfig(level=logging.WARNING)
|
||||||
|
|
||||||
query = sys.argv[1] if len(sys.argv) > 1 else "Hva er OPAX A+ mode?"
|
query = sys.argv[1] if len(sys.argv) > 1 else "Hva er OPAX A+ mode?"
|
||||||
requested_mode = sys.argv[2] if len(sys.argv) > 2 else "A"
|
requested_mode = sys.argv[2] if len(sys.argv) > 2 else "A"
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue
Block a user