OSVauco/agents/core-logic/app.py

# agents/core-logic/app.py
# OSVauco-NMTMD-GCOS — FastAPI HTTP entrypoint for Cloud Run

import os
import sys
import logging
from contextlib import asynccontextmanager
from pathlib import Path

from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import JSONResponse, FileResponse
from pydantic import BaseModel
import vertexai

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', '..'))

try:
    from google.adk.runners import Runner
    from google.adk.sessions import InMemorySessionService
    from google.genai.types import Content, Part
    from agent import build_agent
    from agents.recommendations_engine import get_recommendations
    from token_logger import create_bq_table_if_not_exists
except ImportError as e:
    logger.error(f"Failed to import ADK dependencies: {e}")
    raise

session_service = InMemorySessionService()
root_agent = None
APP_NAME           = os.environ.get("CLOUD_RUN_SERVICE", "gcp-orchestrator")
PROJECT_ID         = os.environ.get("GOOGLE_CLOUD_PROJECT", "propane-will-491900-m5")
BQ_BILLING_DATASET = os.environ.get("BQ_BILLING_DATASET", "billing_data")
CLOUD_BUILD_TRIGGER_ID = os.environ.get("CLOUD_BUILD_TRIGGER_ID", "")

# Static folder: /app/static/ (two levels up from agents/core-logic)
STATIC_DIR = Path(__file__).parent.parent.parent / "static"

_budget_store: dict = {}


@asynccontextmanager
async def lifespan(app: FastAPI):
    global root_agent
    logger.info(f"OSVauco agent '{APP_NAME}' starting up")
    vertexai.init(project=PROJECT_ID, location=os.environ.get("GOOGLE_CLOUD_LOCATION", "us-central1"))
    create_bq_table_if_not_exists()
    root_agent, _, _ = build_agent(mode="light")
    logger.info(f"Static dir: {STATIC_DIR} (exists={STATIC_DIR.exists()})")
    yield
    logger.info(f"OSVauco agent '{APP_NAME}' shutting down")


app = FastAPI(
    title="OSVauco GCP Agent",
    description="ADK-based multi-agent orchestrator on Cloud Run",
    version="1.0.0",
    lifespan=lifespan,
)


class RunRequest(BaseModel):
    user_id: str
    session_id: str
    message: str

class RunResponse(BaseModel):
    user_id: str
    session_id: str
    response: str

class BudgetRequest(BaseModel):
    budget: float
    user: str = "default"


async def _ensure_session(user_id: str, session_id: str):
    try:
        session = await session_service.get_session(
            app_name=APP_NAME, user_id=user_id, session_id=session_id
        )
        if session is not None:
            return session
    except Exception:
        pass
    session = await session_service.create_session(
        app_name=APP_NAME, user_id=user_id, session_id=session_id
    )
    logger.info(f"Created new session: {session_id} for user: {user_id}")
    return session


# ── root — serve opax.html ────────────────────────────────────────────────────
@app.get("/")
async def root():
    index = STATIC_DIR / "opax.html"
    if not index.exists():
        logger.error(f"opax.html not found at {index}")
        raise HTTPException(status_code=404, detail=f"opax.html not found at {index}")
    return FileResponse(str(index), media_type="text/html")


# ── static files ──────────────────────────────────────────────────────────────
@app.get("/static/{filename}")
async def static_file(filename: str):
    filepath = STATIC_DIR / filename
    if not filepath.exists():
        raise HTTPException(status_code=404, detail=f"{filename} not found")
    return FileResponse(str(filepath))


# ── helse ────────────────────────────────────────────────────────────────────
@app.get("/health")
async def health():
    return JSONResponse({"status": "ok", "service": APP_NAME})


# ── agent run ────────────────────────────────────────────────────────────────
@app.post("/run", response_model=RunResponse)
async def run(req: RunRequest):
    try:
        await _ensure_session(req.user_id, req.session_id)
        runner = Runner(agent=root_agent, app_name=APP_NAME, session_service=session_service)
        user_content = Content(role="user", parts=[Part(text=req.message)])
        final_response = ""
        async for event in runner.run_async(
            user_id=req.user_id, session_id=req.session_id, new_message=user_content
        ):
            if event.is_final_response() and event.content:
                for part in event.content.parts:
                    if part.text:
                        final_response += part.text
        logger.info(f"[{req.user_id}/{req.session_id}] Response length: {len(final_response)}")
        return RunResponse(user_id=req.user_id, session_id=req.session_id, response=final_response)
    except Exception as e:
        logger.error(f"Agent run failed: {e}", exc_info=True)
        raise HTTPException(status_code=500, detail=str(e))


# ── OX1a: Cloud Build status ──────────────────────────────────────────────────
@app.get("/opax/build-status")
async def build_status():
    """
    OX1a — Hent siste Cloud Build-kjøring.
    Status: queued|working|success|failure|timeout|cancelled|unknown
    """
    try:
        from google.cloud.devtools import cloudbuild_v1
        client = cloudbuild_v1.CloudBuildClient()
        request = cloudbuild_v1.ListBuildsRequest(
            project_id=PROJECT_ID,
            filter='trigger_id!=""',
            page_size=5,
        )
        builds = list(client.list_builds(request=request))
        if not builds:
            return JSONResponse({"status": "unknown", "message": "Ingen builds funnet"})
        b = builds[0]
        status_map = {1:"queued",2:"working",3:"success",4:"failure",5:"internal_error",6:"timeout",7:"cancelled"}
        status_str = status_map.get(int(b.status), "unknown")
        duration_s = None
        if b.start_time and b.finish_time:
            duration_s = int(b.finish_time.seconds - b.start_time.seconds)
        elif b.start_time:
            import time
            duration_s = int(time.time() - b.start_time.seconds)
        return JSONResponse({
            "status":      status_str,
            "build_id":    b.id,
            "trigger_id":  b.build_trigger_id or "",
            "branch":      (b.substitutions or {}).get("BRANCH_NAME", "main"),
            "commit":      (b.substitutions or {}).get("SHORT_SHA", ""),
            "duration_s":  duration_s,
            "start_time":  str(b.start_time) if b.start_time else None,
            "finish_time": str(b.finish_time) if b.finish_time else None,
            "log_url":     b.log_url or "",
        })
    except Exception as e:
        logger.error(f"[/opax/build-status] Failed: {e}", exc_info=True)
        return JSONResponse({"status": "error", "message": str(e)}, status_code=200)


# ── CG4a: Token intelligence — per modul ─────────────────────────────────────
@app.get("/billing/tokens/by-module")
async def tokens_by_module(days: int = 30):
    """
    CG4a — Token-forbruk aggregert per module_name + caller_type.
    Brukes av opax.html TUI + fremtidig token-panel i billing dashboard.
    Fallback: tom liste hvis BQ mangler data (OQ-15 safe).
    """
    try:
        from google.cloud import bigquery
        client = bigquery.Client(project=PROJECT_ID)
        query = f"""
            SELECT
              COALESCE(module_name, 'ukjent')   AS module_name,
              COALESCE(caller_type, 'ukjent')   AS caller_type,
              model_name,
              COUNT(*)                          AS calls,
              SUM(input_tokens)                 AS input_tokens,
              SUM(output_tokens)                AS output_tokens,
              SUM(total_tokens)                 AS total_tokens,
              ROUND(SUM(estimated_cost_usd), 6) AS cost_usd
            FROM `{PROJECT_ID}.{BQ_BILLING_DATASET}.llm_token_usage`
            WHERE timestamp >= TIMESTAMP_SUB(CURRENT_TIMESTAMP(), INTERVAL {int(days)} DAY)
            GROUP BY module_name, caller_type, model_name
            ORDER BY cost_usd DESC
            LIMIT 100
        """
        rows = [
            {
                "module_name":  row.module_name,
                "caller_type":  row.caller_type,
                "model_name":   row.model_name,
                "calls":        row.calls,
                "input_tokens": row.input_tokens,
                "output_tokens":row.output_tokens,
                "total_tokens": row.total_tokens,
                "cost_usd":     float(row.cost_usd),
            }
            for row in client.query(query).result()
        ]
        return JSONResponse({"period_days": days, "rows": rows})
    except Exception as e:
        logger.error(f"[/billing/tokens/by-module] Failed: {e}", exc_info=True)
        return JSONResponse({"period_days": days, "rows": [], "warning": str(e)})


# ── CG4c: Prosjektestimator ───────────────────────────────────────────────────
@app.get("/billing/tokens/estimate")
async def tokens_estimate(complexity: str = "medium"):
    """
    CG4c — Prosjektestimator.
    Henter snitt-tokens per kall siste 30 dager, multipliserer med kompleksitetsfaktor.
    complexity: low | medium | high | extreme
    Returnerer estimert kostnad i USD for et nytt oppdrag.
    """
    factors = {"low": 0.5, "medium": 1.0, "high": 2.5, "extreme": 6.0}
    factor = factors.get(complexity, 1.0)
    try:
        from google.cloud import bigquery
        client = bigquery.Client(project=PROJECT_ID)
        query = f"""
            SELECT
              AVG(total_tokens)                 AS avg_tokens_per_call,
              AVG(estimated_cost_usd)           AS avg_cost_per_call,
              COUNT(*)                          AS total_calls,
              SUM(estimated_cost_usd)           AS total_cost_usd
            FROM `{PROJECT_ID}.{BQ_BILLING_DATASET}.llm_token_usage`
            WHERE timestamp >= TIMESTAMP_SUB(CURRENT_TIMESTAMP(), INTERVAL 30 DAY)
        """
        result = list(client.query(query).result())
        if not result or result[0].avg_tokens_per_call is None:
            return JSONResponse({"complexity": complexity, "factor": factor, "estimated_usd": None, "warning": "Ingen data ennå (OQ-15)"})
        row = result[0]
        avg_cost = float(row.avg_cost_per_call or 0)
        CALLS_PER_TASK = 20
        estimated = round(avg_cost * CALLS_PER_TASK * factor, 4)
        return JSONResponse({
            "complexity":        complexity,
            "factor":            factor,
            "avg_tokens_per_call": round(float(row.avg_tokens_per_call or 0), 1),
            "avg_cost_per_call": round(avg_cost, 6),
            "total_calls_30d":   row.total_calls,
            "total_cost_30d_usd": round(float(row.total_cost_usd or 0), 4),
            "estimated_usd":     estimated,
            "calls_assumed":     CALLS_PER_TASK,
        })
    except Exception as e:
        logger.error(f"[/billing/tokens/estimate] Failed: {e}", exc_info=True)
        return JSONResponse({"complexity": complexity, "estimated_usd": None, "warning": str(e)})


# ── billing: tokens summary (eksisterende) ────────────────────────────────────
@app.get("/billing/tokens/summary")
async def billing_tokens_summary():
    try:
        from google.cloud import bigquery
        client = bigquery.Client(project=PROJECT_ID)
        query = f"""
            SELECT
              module_name, caller_type, model_name,
              SUM(input_tokens)       AS total_input_tokens,
              SUM(output_tokens)      AS total_output_tokens,
              SUM(total_tokens)       AS total_tokens,
              SUM(estimated_cost_usd) AS total_cost_usd
            FROM `{PROJECT_ID}.{BQ_BILLING_DATASET}.llm_token_usage`
            WHERE timestamp >= TIMESTAMP_SUB(CURRENT_TIMESTAMP(), INTERVAL 30 DAY)
            GROUP BY module_name, caller_type, model_name
            ORDER BY total_cost_usd DESC
        """
        rows = [
            {
                "module_name":         row.module_name,
                "caller_type":         row.caller_type,
                "model_name":          row.model_name,
                "total_input_tokens":  row.total_input_tokens,
                "total_output_tokens": row.total_output_tokens,
                "total_tokens":        row.total_tokens,
                "total_cost_usd":      round(float(row.total_cost_usd), 6),
            }
            for row in client.query(query).result()
        ]
        return JSONResponse({"period_days": 30, "rows": rows})
    except Exception as e:
        logger.error(f"[/billing/tokens/summary] Failed: {e}", exc_info=True)
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/billing/recommendations")
async def billing_recommendations(budget: float = 500.0):
    try:
        recommendations = get_recommendations(budget)
        return JSONResponse(recommendations)
    except Exception as e:
        logger.error(f"[/billing/recommendations] Failed: {e}", exc_info=True)
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/billing/by-service")
async def billing_by_service(days: int = 30):
    try:
        from ml.billing_agent import BillingAgent
        return JSONResponse(BillingAgent().get_service_totals(days))
    except Exception as e:
        logger.error(f"[/billing/by-service] Failed: {e}", exc_info=True)
        return JSONResponse([])


@app.get("/billing/anomalies")
async def billing_anomalies():
    try:
        from ml.billing_agent import BillingAgent
        return JSONResponse(BillingAgent().get_anomalies())
    except Exception as e:
        logger.error(f"[/billing/anomalies] Failed: {e}", exc_info=True)
        return JSONResponse({"anomalies": []})


@app.get("/billing/history")
async def billing_history(days: int = 30):
    try:
        from ml.billing_agent import BillingAgent
        return JSONResponse({"history": BillingAgent().get_daily_history(days)})
    except Exception as e:
        logger.error(f"[/billing/history] Failed: {e}", exc_info=True)
        return JSONResponse({"history": []})


@app.get("/billing/summary")
async def billing_summary():
    try:
        from ml.billing_agent import BillingAgent
        return JSONResponse(BillingAgent().get_summary())
    except Exception as e:
        logger.error(f"[/billing/summary] Failed: {e}", exc_info=True)
        return JSONResponse({"summary": []})


@app.get("/billing/live")
async def billing_live():
    try:
        from ml.billing_agent import BillingAgent
        return JSONResponse(BillingAgent().get_forecast())
    except Exception as e:
        logger.error(f"[/billing/live] Failed: {e}", exc_info=True)
        return JSONResponse({})


@app.get("/billing/budget")
async def get_budget(user: str = "default"):
    return JSONResponse({"budget": _budget_store.get(user, 500.0), "user": user})


@app.post("/billing/budget")
async def set_budget(req: BudgetRequest):
    if req.budget <= 0:
        raise HTTPException(status_code=400, detail="Budget must be > 0")
    _budget_store[req.user] = req.budget
    return JSONResponse({"budget": req.budget, "user": req.user})