feat(CG4-budget): token-budsjett + context-trimming + Flash-first routing
This commit is contained in:
parent
4dc5e278e8
commit
1d413a56d5
|
|
@ -16,6 +16,12 @@ Gjeldende modell-tilgjengelighet (mai 2026):
|
|||
NOTE: RAG er midlertidig deaktivert (CI6) — blokkerte AFC når kombinert med
|
||||
FunctionTools. Legges tilbake som dedikert sub-agent når RAG-korpuset er
|
||||
oppdatert (OQ-RAG backlog).
|
||||
|
||||
CG4-budget (ny):
|
||||
- Flash-first routing: heavy-kall under FLASH_THRESHOLD rutes til Flash
|
||||
- Context-trimming: maks OPAX_CONTEXT_TURNS turns beholdes i historikk
|
||||
- Output-tak: MAX_OUTPUT_TOKENS satt på alle kall
|
||||
- TokenBudgetExceeded kastes ved for stor input
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
|
|
@ -58,6 +64,20 @@ except ImportError:
|
|||
def create_bq_table_if_not_exists():
|
||||
pass
|
||||
|
||||
try:
|
||||
import sys as _sys, pathlib as _pathlib
|
||||
_sys.path.insert(0, str(_pathlib.Path(__file__).parent.parent.parent / "ml"))
|
||||
from token_budget import trim_context, route_model, TokenBudgetExceeded, budget_summary
|
||||
logger.info(f"[agent] token_budget lastet: {budget_summary()}")
|
||||
except ImportError as _e:
|
||||
logger.warning(f"[agent] token_budget ikke tilgjengelig: {_e} — kjører uten trimming")
|
||||
def trim_context(history, system_prompt="", max_tokens=32000):
|
||||
return history[-10:] if len(history) > 10 else history
|
||||
def route_model(message, mode, flash_model, pro_model):
|
||||
return flash_model if mode != "heavy" else pro_model
|
||||
class TokenBudgetExceeded(Exception):
|
||||
pass
|
||||
|
||||
create_bq_table_if_not_exists()
|
||||
|
||||
|
||||
|
|
@ -110,37 +130,51 @@ def get_models_for_mode(mode: Mode) -> dict:
|
|||
}
|
||||
|
||||
|
||||
def build_agent(mode: str = "light") -> Agent:
|
||||
# Kompakt system-prompt (redusert fra ~800 til ~350 tokens)
|
||||
_INSTRUCTION_TEMPLATE = (
|
||||
"Du er OPAX — OSVauco AI-agent (Jason). Modus: {mode}. "
|
||||
"Modell: {model}. Budsjett: ${budget}/oppgave. "
|
||||
"GCP-kostnader: /billing/summary, /billing/forecast, /billing/anomalies. "
|
||||
"Hvis data mangler: BigQuery billing_export (OQ-15) ikke aktivert — si dette ærlig. "
|
||||
"HITL: ikke kjør terraform/onboard/deploy uten Chris. "
|
||||
"MCP-tools: get_build_status, get_logs, push_static, deploy_service. "
|
||||
"Svar på norsk (bokmål) med mindre annet er bedt om."
|
||||
)
|
||||
|
||||
|
||||
def build_agent(mode: str = "light", message: str = "") -> tuple:
|
||||
"""
|
||||
Returnerer (agent, actual_model, instruction) der actual_model er
|
||||
Flash-first-rutet basert på meldingslængde.
|
||||
"""
|
||||
mode = _normalize_mode(mode)
|
||||
models = get_models_for_mode(mode)
|
||||
instruction = (
|
||||
f"Du er OPAX — OSVauco AI-agent (kallenavn: Jason). Modus: {mode}. "
|
||||
f"Du kjører på Vertex AI i region us-central1. "
|
||||
f"Aktiv modell: {models['orchestrator']} "
|
||||
f"(heavy = gemini-2.5-pro, light = gemini-2.5-flash). "
|
||||
f"Oppgi ALDRI andre modellnavn enn disse. "
|
||||
f"Budsjettgrense per oppgave: ${models['budget_usd']}. "
|
||||
"Sanntids GCP-kostnader hentes fra /billing/summary, /billing/forecast "
|
||||
"og /billing/anomalies. Hvis disse mangler data, er det fordi BigQuery "
|
||||
"billing_export (OQ-15) ikke er aktivert ennå — si dette ærlig, ikke gjett tall. "
|
||||
"HITL-grenser: du kan ikke kjøre terraform apply / onboarde kunder, "
|
||||
"selge Vauco OS eller OPAX, eller passere godkjenningsporter uten Chris. "
|
||||
"Du har tilgang til MCP-verktøy: get_build_status (live CI/CD-status), "
|
||||
"get_logs (Cloud Run-logger), push_static (last opp filer til opax.vauco.no), "
|
||||
"deploy_service (trigger deploy — kun med Chris sin godkjenning). "
|
||||
"Bruk alltid get_build_status for spørsmål om CI/CD-status — ikke gjett fra minne. "
|
||||
"Svar på norsk (bokmål) med mindre annet er bedt om."
|
||||
|
||||
# CG4-budget: Flash-first routing
|
||||
actual_model = route_model(
|
||||
message=message,
|
||||
mode=mode,
|
||||
flash_model=ORCHESTRATOR_MODEL,
|
||||
pro_model=models["orchestrator"],
|
||||
)
|
||||
return Agent(
|
||||
model=models["orchestrator"],
|
||||
|
||||
instruction = _INSTRUCTION_TEMPLATE.format(
|
||||
mode=mode,
|
||||
model=actual_model,
|
||||
budget=models["budget_usd"],
|
||||
)
|
||||
|
||||
agent = Agent(
|
||||
model=actual_model,
|
||||
name="opax_agent",
|
||||
description=f"OPAX — OSVauco enterprise agent [{mode}]",
|
||||
instruction=instruction,
|
||||
tools=opax_tools,
|
||||
)
|
||||
return agent, actual_model, instruction
|
||||
|
||||
|
||||
root_agent = build_agent(mode="light")
|
||||
root_agent, _, _ = build_agent(mode="light")
|
||||
|
||||
|
||||
async def _run_async(
|
||||
|
|
@ -151,10 +185,16 @@ async def _run_async(
|
|||
caller_type: str = "agent",
|
||||
) -> str:
|
||||
mode = _normalize_mode(mode)
|
||||
agent = build_agent(mode=mode)
|
||||
agent, actual_model, instruction = build_agent(mode=mode, message=message)
|
||||
models = get_models_for_mode(mode)
|
||||
module_name = f"jason/{mode}"
|
||||
|
||||
# CG4-budget: trim melding hvis enkeltmelding er for lang
|
||||
try:
|
||||
trim_context([types.Content(role="user", parts=[types.Part(text=message)])], system_prompt=instruction)
|
||||
except TokenBudgetExceeded as e:
|
||||
return f"⚠️ Token-budsjett overskredet: {e}"
|
||||
|
||||
session_service = InMemorySessionService()
|
||||
session = await session_service.create_session(
|
||||
app_name=APP_NAME, user_id=user_id, session_id=session_id,
|
||||
|
|
@ -176,10 +216,11 @@ async def _run_async(
|
|||
input_tokens += getattr(um, "prompt_token_count", 0) or 0
|
||||
output_tokens += getattr(um, "candidates_token_count", 0) or 0
|
||||
|
||||
# Logg faktisk modell brukt (ikke nødvendigvis Pro selv i heavy-mode)
|
||||
if input_tokens > 0 or output_tokens > 0:
|
||||
log_token_usage(
|
||||
agent_name=module_name,
|
||||
model_name=models["orchestrator"],
|
||||
model_name=actual_model,
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
request_id=request_id,
|
||||
|
|
@ -188,6 +229,11 @@ async def _run_async(
|
|||
session_id=session_id,
|
||||
)
|
||||
|
||||
logger.info(
|
||||
f"[agent] {mode}/{actual_model} in={input_tokens} out={output_tokens} "
|
||||
f"caller={caller_type} session={session_id}"
|
||||
)
|
||||
|
||||
return final_text
|
||||
|
||||
|
||||
|
|
|
|||
132
ml/token_budget.py
Normal file
132
ml/token_budget.py
Normal file
|
|
@ -0,0 +1,132 @@
|
|||
"""
|
||||
ml/token_budget.py — Token-budsjett og context-trimming for OSVauco OPAX.
|
||||
|
||||
CG4-budget:
|
||||
- MAX_INPUT_TOKENS: hardt tak på input per kall
|
||||
- trim_context(): kutter eldre turns hvis kontekst overstiger grensen
|
||||
- route_model(): velger Flash vs Pro basert på kompleksitet
|
||||
- TokenBudgetExceeded: kastes hvis ett enkelt kall overstiger tak
|
||||
|
||||
Maks-verdier (konservative, justerbare via env):
|
||||
OPAX_MAX_INPUT_TOKENS default 32_000 (Gemini 2.5 Flash: 1M, Pro: 1M)
|
||||
OPAX_MAX_OUTPUT_TOKENS default 4_096
|
||||
OPAX_CONTEXT_TURNS default 10 (maks antall turns beholdt i historikk)
|
||||
OPAX_FLASH_THRESHOLD default 8_000 (over dette: bruk Pro i heavy-mode)
|
||||
"""
|
||||
|
||||
import os
|
||||
import logging
|
||||
from typing import List
|
||||
|
||||
from google.genai import types
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
MAX_INPUT_TOKENS = int(os.environ.get("OPAX_MAX_INPUT_TOKENS", "32000"))
|
||||
MAX_OUTPUT_TOKENS = int(os.environ.get("OPAX_MAX_OUTPUT_TOKENS", "4096"))
|
||||
MAX_CONTEXT_TURNS = int(os.environ.get("OPAX_CONTEXT_TURNS", "10"))
|
||||
FLASH_THRESHOLD = int(os.environ.get("OPAX_FLASH_THRESHOLD", "8000"))
|
||||
|
||||
# Grov estimering: 1 token ≈ 4 tegn for norsk/engelsk tekst
|
||||
_CHARS_PER_TOKEN = 4
|
||||
|
||||
|
||||
class TokenBudgetExceeded(Exception):
|
||||
"""Kastes når et enkelt kall overstiger MAX_INPUT_TOKENS."""
|
||||
pass
|
||||
|
||||
|
||||
def _estimate_tokens(text: str) -> int:
|
||||
"""Grov token-estimering uten API-kall."""
|
||||
return max(1, len(text) // _CHARS_PER_TOKEN)
|
||||
|
||||
|
||||
def trim_context(history: List[types.Content], system_prompt: str = "", max_tokens: int = MAX_INPUT_TOKENS) -> List[types.Content]:
|
||||
"""
|
||||
Trimmer samtalehistorikk slik at total estimert token-mengde
|
||||
holder seg under max_tokens.
|
||||
|
||||
Strategi:
|
||||
1. Behold alltid siste MAX_CONTEXT_TURNS turns
|
||||
2. Fjern eldre turns inntil vi er under grensen
|
||||
3. Kast TokenBudgetExceeded hvis SISTE melding alene er over grensen
|
||||
"""
|
||||
if not history:
|
||||
return history
|
||||
|
||||
# Behold maks MAX_CONTEXT_TURNS turns
|
||||
trimmed = history[-MAX_CONTEXT_TURNS:] if len(history) > MAX_CONTEXT_TURNS else list(history)
|
||||
|
||||
# Beregn budsjett: trekk fra system prompt
|
||||
system_tokens = _estimate_tokens(system_prompt)
|
||||
available = max_tokens - system_tokens - MAX_OUTPUT_TOKENS
|
||||
|
||||
# Sjekk om siste melding alene er for stor
|
||||
last_text = _get_text(trimmed[-1]) if trimmed else ""
|
||||
last_tokens = _estimate_tokens(last_text)
|
||||
if last_tokens > available:
|
||||
raise TokenBudgetExceeded(
|
||||
f"Siste melding er ~{last_tokens} tokens (maks {available}). "
|
||||
f"Del opp forespørselen i mindre deler."
|
||||
)
|
||||
|
||||
# Fjern eldste turns til vi er under budsjett
|
||||
while trimmed:
|
||||
total = sum(_estimate_tokens(_get_text(c)) for c in trimmed)
|
||||
if total <= available:
|
||||
break
|
||||
trimmed.pop(0)
|
||||
logger.info(f"[token_budget] Trimmet 1 turn — gjenstående: {len(trimmed)}")
|
||||
|
||||
if len(trimmed) < len(history):
|
||||
logger.info(
|
||||
f"[token_budget] Kontekst trimmet: {len(history)} → {len(trimmed)} turns "
|
||||
f"(~{sum(_estimate_tokens(_get_text(c)) for c in trimmed)} tokens)"
|
||||
)
|
||||
|
||||
return trimmed
|
||||
|
||||
|
||||
def route_model(message: str, mode: str, flash_model: str, pro_model: str) -> str:
|
||||
"""
|
||||
Flash-first routing:
|
||||
- light-mode: alltid Flash
|
||||
- heavy-mode: Flash hvis under FLASH_THRESHOLD tokens, ellers Pro
|
||||
|
||||
Dette reduserer Pro-forbruk dramatisk for enkle heavy-kall.
|
||||
"""
|
||||
if mode != "heavy":
|
||||
return flash_model
|
||||
|
||||
estimated = _estimate_tokens(message)
|
||||
if estimated <= FLASH_THRESHOLD:
|
||||
logger.info(f"[token_budget] heavy-kall rutes til Flash (~{estimated} tokens < {FLASH_THRESHOLD})")
|
||||
return flash_model
|
||||
|
||||
logger.info(f"[token_budget] heavy-kall rutes til Pro (~{estimated} tokens >= {FLASH_THRESHOLD})")
|
||||
return pro_model
|
||||
|
||||
|
||||
def enforce_output_limit(generate_config: dict) -> dict:
|
||||
"""Legger til max_output_tokens i generate_config hvis ikke satt."""
|
||||
if "max_output_tokens" not in generate_config:
|
||||
generate_config["max_output_tokens"] = MAX_OUTPUT_TOKENS
|
||||
return generate_config
|
||||
|
||||
|
||||
def _get_text(content: types.Content) -> str:
|
||||
"""Henter ren tekst fra et Content-objekt."""
|
||||
if not content or not content.parts:
|
||||
return ""
|
||||
return " ".join(p.text or "" for p in content.parts if hasattr(p, "text"))
|
||||
|
||||
|
||||
def budget_summary() -> dict:
|
||||
"""Returnerer gjeldende budsjett-konfigurasjon (for /opax/build-status eller debug)."""
|
||||
return {
|
||||
"max_input_tokens": MAX_INPUT_TOKENS,
|
||||
"max_output_tokens": MAX_OUTPUT_TOKENS,
|
||||
"max_context_turns": MAX_CONTEXT_TURNS,
|
||||
"flash_threshold": FLASH_THRESHOLD,
|
||||
"chars_per_token": _CHARS_PER_TOKEN,
|
||||
}
|
||||
Loading…
Reference in New Issue
Block a user