fix: inline opax_mcp_client
This commit is contained in:
parent
c935841011
commit
afabe8d2a9
|
|
@ -1,228 +1 @@
|
|||
#!/usr/bin/env python3
|
||||
"""
|
||||
agent.py — OSVauco OPAX agent (Jason Vauger).
|
||||
|
||||
Modes:
|
||||
light — gemini-2.5-flash, $1/task
|
||||
heavy — gemini-2.5-pro, $3/task
|
||||
|
||||
Legacy 'A' / 'A+' normaliseres automatisk.
|
||||
Autoriserte brukere for heavy: opax, admin
|
||||
|
||||
NOTE: opax_mcp_client.py ligg i same mappe (core-logic/) for Vertex AI-kompatibilitet.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import logging
|
||||
import uuid
|
||||
from typing import Literal
|
||||
|
||||
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__)
|
||||
|
||||
PROJECT_ID = os.environ.get("GOOGLE_CLOUD_PROJECT", "propane-will-491900-m5")
|
||||
LOCATION = os.environ.get("GOOGLE_CLOUD_LOCATION", "global")
|
||||
|
||||
ORCHESTRATOR_MODEL = os.environ.get("ORCHESTRATOR_MODEL", "gemini-2.5-flash")
|
||||
SUBAGENT_MODEL = os.environ.get("SUBAGENT_MODEL", "gemini-2.5-flash")
|
||||
REASONING_MODEL = os.environ.get("REASONING_MODEL", "gemini-2.5-flash")
|
||||
|
||||
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_REASONING = os.environ.get("HEAVY_REASONING_MODEL", "gemini-2.5-flash")
|
||||
|
||||
BUDGET_LIGHT = float(os.environ.get("BUDGET_A_USD_PER_TASK", "1.0"))
|
||||
BUDGET_HEAVY = float(os.environ.get("HEAVY_MODE_BUDGET_USD_PER_DAY", "3.0"))
|
||||
|
||||
HEAVY_MODE_ALLOWED_USERS = ["opax", "admin"]
|
||||
|
||||
Mode = Literal["light", "heavy"]
|
||||
APP_NAME = "opax"
|
||||
|
||||
try:
|
||||
from token_logger import log_token_usage, create_bq_table_if_not_exists
|
||||
except ImportError:
|
||||
def log_token_usage(*args, **kwargs): pass
|
||||
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}")
|
||||
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()
|
||||
|
||||
|
||||
def _normalize_mode(mode: str) -> str:
|
||||
mapping = {"A": "light", "A+": "heavy", "light": "light", "heavy": "heavy"}
|
||||
if mode not in mapping:
|
||||
raise ValueError(f"Invalid mode '{mode}'. Must be 'light' or 'heavy'.")
|
||||
return mapping[mode]
|
||||
|
||||
|
||||
def authorize_mode(user_id: str, mode: str) -> None:
|
||||
mode = _normalize_mode(mode)
|
||||
if mode == "heavy" and user_id not in HEAVY_MODE_ALLOWED_USERS:
|
||||
raise PermissionError(f"User '{user_id}' is not authorized for heavy mode.")
|
||||
|
||||
|
||||
# —— OPAX-MCP FunctionTools ——
|
||||
# opax_mcp_client.py er i same mappe — Vertex AI finn den alltid
|
||||
opax_tools = []
|
||||
try:
|
||||
from opax_mcp_client import get_all_function_tools
|
||||
opax_tools = get_all_function_tools()
|
||||
logger.info(f"OPAX-MCP tools lastet: {len(opax_tools)} tools")
|
||||
except Exception as e:
|
||||
logger.warning(f"OPAX-MCP tools ikkje tilgjengeleg: {e}")
|
||||
|
||||
|
||||
def get_models_for_mode(mode: Mode) -> dict:
|
||||
if mode == "heavy":
|
||||
return {
|
||||
"orchestrator": HEAVY_ORCHESTRATOR,
|
||||
"subagent": HEAVY_SUBAGENT,
|
||||
"reasoning": HEAVY_REASONING,
|
||||
"budget_usd": BUDGET_HEAVY,
|
||||
}
|
||||
return {
|
||||
"orchestrator": ORCHESTRATOR_MODEL,
|
||||
"subagent": SUBAGENT_MODEL,
|
||||
"reasoning": REASONING_MODEL,
|
||||
"budget_usd": BUDGET_LIGHT,
|
||||
}
|
||||
|
||||
|
||||
_INSTRUCTION_TEMPLATE = (
|
||||
"Du er Jason Vauger — OPAX-agenten for Vauco AS. Modus: {mode}. Modell: {model}. "
|
||||
"Du har tilgang til alle OPAX-tools via opax-mcp. "
|
||||
"Tools du kan bruke: list_commits, get_file, push_file, create_issue, list_open_issues (Gitea), "
|
||||
"get_billing_summary, get_billing_credits, get_billing_anomalies, "
|
||||
"get_build_status, get_state, get_telemetry, run_terminal, "
|
||||
"create_invite, list_customers, send_webhook, send_email, get_notify_channels. "
|
||||
"Når brukaren ber om data — KALL ALLTID rett tool før du svarar. "
|
||||
"Svar på norsk (bokmål) med mindre annet er bedt om. "
|
||||
"HITL: ikkje kjør terraform/deploy utan godkjenning frå Chris."
|
||||
)
|
||||
|
||||
|
||||
def build_agent(mode: str = "light", message: str = "") -> tuple:
|
||||
mode = _normalize_mode(mode)
|
||||
models = get_models_for_mode(mode)
|
||||
actual_model = route_model(
|
||||
message=message,
|
||||
mode=mode,
|
||||
flash_model=ORCHESTRATOR_MODEL,
|
||||
pro_model=models["orchestrator"],
|
||||
)
|
||||
instruction = _INSTRUCTION_TEMPLATE.format(mode=mode, model=actual_model)
|
||||
agent = Agent(
|
||||
model=actual_model,
|
||||
name="jason_vauger",
|
||||
description=f"Jason Vauger — OPAX enterprise agent [{mode}]",
|
||||
instruction=instruction,
|
||||
tools=opax_tools,
|
||||
)
|
||||
return agent, actual_model, instruction
|
||||
|
||||
|
||||
root_agent, _, _ = build_agent(mode="light")
|
||||
|
||||
|
||||
async def _run_async(
|
||||
message: str,
|
||||
user_id: str,
|
||||
session_id: str,
|
||||
mode: str,
|
||||
caller_type: str = "agent",
|
||||
) -> str:
|
||||
mode = _normalize_mode(mode)
|
||||
agent, actual_model, instruction = build_agent(mode=mode, message=message)
|
||||
models = get_models_for_mode(mode)
|
||||
module_name = f"jason/{mode}"
|
||||
|
||||
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,
|
||||
)
|
||||
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 = ""
|
||||
input_tokens = 0
|
||||
output_tokens = 0
|
||||
request_id = str(uuid.uuid4())
|
||||
|
||||
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 ""
|
||||
if hasattr(event, "usage_metadata") and event.usage_metadata:
|
||||
um = event.usage_metadata
|
||||
input_tokens += getattr(um, "prompt_token_count", 0) or 0
|
||||
output_tokens += getattr(um, "candidates_token_count", 0) or 0
|
||||
|
||||
if input_tokens > 0 or output_tokens > 0:
|
||||
log_token_usage(
|
||||
agent_name=module_name,
|
||||
model_name=actual_model,
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
request_id=request_id,
|
||||
module_name=module_name,
|
||||
caller_type=caller_type,
|
||||
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
|
||||
|
||||
|
||||
def run(
|
||||
message: str,
|
||||
user_id: str = "opax",
|
||||
session_id: str = "default",
|
||||
mode: str = "light",
|
||||
caller_type: str = "agent",
|
||||
) -> str:
|
||||
mode = _normalize_mode(mode)
|
||||
authorize_mode(user_id, mode)
|
||||
return asyncio.run(_run_async(
|
||||
message=message,
|
||||
user_id=user_id,
|
||||
session_id=session_id,
|
||||
mode=mode,
|
||||
caller_type=caller_type,
|
||||
))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
logging.basicConfig(level=logging.WARNING)
|
||||
query = sys.argv[1] if len(sys.argv) > 1 else "Hva er OPAX?"
|
||||
mode = sys.argv[2] if len(sys.argv) > 2 else "light"
|
||||
caller_type = sys.argv[3] if len(sys.argv) > 3 else "cli"
|
||||
print(f"Mode: {mode} | Query: {query}")
|
||||
print("-" * 60)
|
||||
print(run(message=query, user_id="opax", mode=mode, caller_type=caller_type))
|
||||
404: Not Found
|
||||
Loading…
Reference in New Issue
Block a user