app.py

#!/usr/bin/env python3
"""
Redis LangCache demo - FastAPI backend.

Serves the static UI (static/) and a small JSON API that wraps LangCache + OpenAI:
  POST /api/login   -> password gate, sets an HttpOnly session cookie
  GET  /api/meta    -> service metadata + global counters (observability panel)
  POST /api/chat    -> semantic-cache-first answer with full telemetry
  POST /api/flush   -> delete cache entries by scope (index is never dropped)
  GET  /healthz     -> liveness probe (no auth)

Demo semantics (ported from legacy/main_demo_released.py):
  - Scoped caching by company / business unit / person via LangCache attributes.
  - Identity handling: "My role is X" (EN) or "Minha função é X" (PT-BR) is
    stored under an exact key per scope; the user's name is never cached.
  - Ambiguous prompts ("deploy", "pipeline", ...) are rewritten with the domain
    inferred from the persona's business unit before hitting cache/LLM.
"""

import hmac
import json
import os
import re
import secrets
import time
from collections import deque
from typing import Any, Dict, List, Optional, Tuple

import uvicorn
from dotenv import load_dotenv
from fastapi import FastAPI, Request, Response
from fastapi.responses import FileResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from openai import OpenAI
from pydantic import BaseModel

try:
    from langcache import LangCache
    from langcache.models import SearchStrategy
    import langcache as langcache_pkg
except Exception:  # pragma: no cover - demo still boots without the SDK
    LangCache = None
    SearchStrategy = None
    langcache_pkg = None

# ============== Env & clients ==============
load_dotenv()

LANGCACHE_API_KEY = os.getenv("LANGCACHE_API_KEY") or os.getenv("LANGCACHE_SERVICE_KEY")
LANGCACHE_CACHE_ID = os.getenv("LANGCACHE_CACHE_ID", "")
LANGCACHE_BASE_URL = os.getenv("LANGCACHE_BASE_URL", "https://gcp-us-east4.langcache.redis.io")

OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
OPENAI_MODEL = os.getenv("OPENAI_MODEL", "gpt-4o-mini")
if not OPENAI_API_KEY:
    raise SystemExit("OPENAI_API_KEY is required.")
openai_client = OpenAI(api_key=OPENAI_API_KEY)

APP_PASSWORD = os.getenv("APP_PASSWORD", "secret42")
SESSION_TOKEN = secrets.token_hex(32)  # rotates on restart; fine for a demo
COOKIE_NAME = "lc_session"

lang_cache: Optional["LangCache"] = None
if LANGCACHE_API_KEY and LANGCACHE_CACHE_ID and LangCache is not None:
    lang_cache = LangCache(
        server_url=LANGCACHE_BASE_URL,
        cache_id=LANGCACHE_CACHE_ID,
        api_key=LANGCACHE_API_KEY,
    )
else:
    print("[WARN] LangCache is not configured; UI runs without a real cache.")

# ============== Global counters (since boot) ==============
STARTED_AT = time.time()
STATS: Dict[str, Any] = {
    "requests": 0,
    "hits": 0,
    "misses": 0,
    "stores": 0,
    "tokens_saved": 0,
    "usd_saved": 0.0,
    "cache_ms": deque(maxlen=300),
    "llm_ms": deque(maxlen=300),
}


def stats_snapshot() -> Dict[str, Any]:
    cache_ms = list(STATS["cache_ms"])
    llm_ms = list(STATS["llm_ms"])
    avg_cache = sum(cache_ms) / len(cache_ms) if cache_ms else 0.0
    avg_llm = sum(llm_ms) / len(llm_ms) if llm_ms else 0.0
    total = STATS["hits"] + STATS["misses"]
    return {
        "requests": STATS["requests"],
        "hits": STATS["hits"],
        "misses": STATS["misses"],
        "stores": STATS["stores"],
        "hit_rate": (STATS["hits"] / total * 100.0) if total else 0.0,
        "tokens_saved": STATS["tokens_saved"],
        "usd_saved": STATS["usd_saved"],
        "avg_cache_ms": avg_cache,
        "avg_llm_ms": avg_llm,
        "speedup": (avg_llm / avg_cache) if avg_cache > 0 and avg_llm > 0 else None,
        "uptime_s": int(time.time() - STARTED_AT),
    }


# ============== Intent / identity (bilingual EN + PT-BR demo triggers) ==============
KEY_NAME = "[IDENTITY:NAME]"
KEY_ROLE = "[IDENTITY:ROLE]"

NAME_TRIGGERS = [
    # EN
    "what is my name", "what's my name", "whats my name",
    "say my name", "tell me my name", "who am i",
    # PT-BR
    "qual é meu nome", "qual eh meu nome", "qual o meu nome", "qual meu nome",
    "como me chamo", "diga meu nome", "fale meu nome", "quem sou eu",
    "qual é o meu nome", "qual eh o meu nome",
]

ROLE_TRIGGERS = [
    # EN
    "what is my role", "what's my role", "whats my role",
    "what is my job title", "what is my position", "what's my position",
    "what do i do at the company", "my role at the company",
    # PT-BR
    "qual é a minha função", "qual eh a minha funcao", "qual a minha função", "qual a minha funcao",
    "qual é minha função", "qual eh minha funcao",
    "qual é meu cargo", "qual eh meu cargo", "qual meu cargo",
    "qual é a minha posição", "qual a minha posição", "minha posição", "minha posicao",
    "qual é o meu papel", "qual o meu papel", "meu papel",
    "o que eu faço na empresa", "o que eu faco na empresa",
]

ROLE_SET_PATTERNS = [
    # EN
    r"\bmy role at the company is\s+(?P<role>.+)$",
    r"\bmy role is\s+(?P<role>.+)$",
    r"\bmy job title is\s+(?P<role>.+)$",
    r"\bi work as\s+(?P<role>.+)$",
    # PT-BR
    r"\bminha função é\s+(?P<role>.+)$",
    r"\bminha funcao é\s+(?P<role>.+)$",
    r"\bminha funcao eh\s+(?P<role>.+)$",
    r"\bmeu cargo é\s+(?P<role>.+)$",
    r"\btrabalho como\s+(?P<role>.+)$",
]


def is_name_prompt(p: str) -> bool:
    p = (p or "").strip().lower()
    return any(t in p for t in NAME_TRIGGERS)


def is_role_prompt(p: str) -> bool:
    p = (p or "").strip().lower()
    return any(t in p for t in ROLE_TRIGGERS)


def try_extract_role_set(p: str) -> Optional[str]:
    txt = re.sub(r"[.!?]\s*$", "", (p or "").strip())
    for pat in ROLE_SET_PATTERNS:
        m = re.search(pat, txt, flags=re.IGNORECASE)
        if m:
            return re.sub(r"[.!?]\s*$", "", m.group("role").strip())
    return None


def normalize_prompt_for_cache(prompt: str) -> Tuple[str, str]:
    if is_name_prompt(prompt):
        return KEY_NAME, "identity:name"
    if is_role_prompt(prompt):
        return KEY_ROLE, "identity:role"
    return f"[FACT]\n{prompt.strip()}", "fact"


def depersonalize_safe(text: str, person: Optional[str]) -> str:
    if not text or not person:
        return text
    original = text.strip()
    t = original
    for pat in [
        rf"^h(?:i|ello),?\s*{re.escape(person)}\s*!?\s*",
        rf"^your\s+name\s+is\s*{re.escape(person)}[.!]?\s*",
        rf"^ol[áa],\s*{re.escape(person)}\s*!\s*",
        rf"^seu\s+nome\s+é\s*{re.escape(person)}[.!]?\s*",
        rf"^seu\s+nome\s+eh\s*{re.escape(person)}[.!]?\s*",
        rf"^voc[êe]\s+se\s+chama\s*{re.escape(person)}[.!]?\s*",
    ]:
        t = re.sub(pat, "", t, flags=re.IGNORECASE)
    t = t.strip()
    return t if t else original


# ============== Domain inference / disambiguation ==============
AMBIGUOUS_TERMS = [
    r"\bdeploy\b", r"\bpipeline\b", r"\bmodel\b", r"\bnetwork\b", r"\bcell\b", r"\bbank\b",
    r"\bc[eé]lula\b", r"\bbanco\b", r"\brede\b", r"\bmodelo\b",
]


def infer_domain(company: str, bu: str, role: Optional[str] = None) -> str:
    text = f"{company} {bu} {role or ''}".lower()
    if any(k in text for k in ["health", "clinic", "saude", "clínica", "clinica", "medic", "hospital"]):
        return "healthcare"
    if any(k in text for k in ["engineering", "engenharia", "software", "dev", "produto", "ti", "tecnologia", "tech"]):
        return "software engineering"
    if any(k in text for k in ["dados", "data", "bi", "analytics"]):
        return "data"
    if any(k in text for k in ["finan", "banco", "bank", "invest", "asset", "seguro", "insurance"]):
        return "corporate finance"
    if any(k in text for k in ["tourism", "turismo", "eco", "aventura", "hotel", "viagem", "travel"]):
        return "tourism"
    return "the user's general field"


def looks_ambiguous(prompt: str) -> bool:
    p = (prompt or "").lower()
    return any(re.search(pat, p, flags=re.IGNORECASE) for pat in AMBIGUOUS_TERMS)


def rewrite_with_domain(prompt: str, domain_label: str) -> str:
    clean = prompt.strip()
    if not clean.endswith("?"):
        clean += "?"
    return f"{clean} (in the context of {domain_label})"


# ============== LLM ==============
FEW_SHOT: List[Dict[str, str]] = [
    {"role": "user", "content": "What is a deploy? (in the context of software engineering)"},
    {"role": "assistant", "content": "A deploy is the process of releasing a new version of software into production."},
    {"role": "user", "content": "What is a pipeline? (in the context of software engineering)"},
    {"role": "assistant", "content": "A pipeline is an automated sequence of steps to build, test and ship code."},
    {"role": "user", "content": "What is a deploy? (in the context of corporate finance)"},
    {"role": "assistant", "content": "In finance, a deploy can refer to rolling out a new process, system or investment for internal use."},
    {"role": "user", "content": "What is a pipeline? (in the context of sales and finance)"},
    {"role": "assistant", "content": "A pipeline is the list of opportunities or revenue forecasts still in progress."},
    {"role": "user", "content": "Explain what machine learning is."},
    {"role": "assistant", "content": "Machine learning is a branch of AI that enables systems to learn patterns from data and make predictions or decisions without being explicitly programmed."},
    {"role": "user", "content": "O que é aprendizado de máquina?"},
    {"role": "assistant", "content": "Aprendizado de máquina é uma área da IA que permite que sistemas aprendam padrões a partir de dados sem programação explícita."},
]


def call_openai(prompt: str, company: str, bu: str) -> str:
    domain = infer_domain(company, bu)
    system_ctx = (
        "Answer briefly and directly. "
        "Answer in the same language as the question. "
        "Do not mention the user's name unless the question is about their name/identity. "
        f"Main context: {domain}. "
        "If the question is ambiguous (e.g. 'deploy', 'pipeline', 'model'), "
        f"answer ONLY in the sense of {domain} and do NOT mention other meanings."
    )
    msgs = [{"role": "system", "content": system_ctx}, *FEW_SHOT, {"role": "user", "content": prompt}]
    resp = openai_client.chat.completions.create(model=OPENAI_MODEL, messages=msgs)
    return resp.choices[0].message.content.strip()


def estimate_tokens(*texts: str) -> int:
    return max(1, sum(len(t or "") for t in texts) // 4)


def calc_savings(tokens_est: int, price_in: float, price_out: float, frac_in: float) -> float:
    """price_in/price_out are USD per 1M tokens."""
    tokens_in = int(tokens_est * frac_in)
    tokens_out = max(0, tokens_est - tokens_in)
    return (tokens_in / 1_000_000.0) * price_in + (tokens_out / 1_000_000.0) * price_out


# ============== Scope attributes ==============
def build_attributes(company: str, bu: str, person: str, isolation: str) -> Dict[str, str]:
    if isolation == "company":
        return {"company": company}
    if isolation == "company+bu":
        return {"company": company, "business_unit": bu}
    if isolation == "company+bu+person":
        return {"company": company, "business_unit": bu, "person": person}
    return {}


# ============== API models ==============
class LoginBody(BaseModel):
    password: str


class ChatBody(BaseModel):
    side: str = "A"
    company: str = "RedisLabs"
    bu: str = "Software-Engineering"
    person: str = "Gabs"
    prompt: str
    isolation: str = "company+bu+person"
    similarity_threshold: float = 0.85  # < 0 disables the threshold
    exact_then_semantic: bool = True
    ttl_seconds: int = 0
    price_in_per_m: float = 0.15
    price_out_per_m: float = 0.60
    frac_in: float = 0.5


class FlushBody(BaseModel):
    company: str = ""
    bu: str = ""
    person: str = ""
    isolation: str = "company+bu+person"


# ============== App ==============
app = FastAPI(title="Redis LangCache Demo")
STATIC_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "static")


def is_authed(request: Request) -> bool:
    cookie = request.cookies.get(COOKIE_NAME, "")
    return hmac.compare_digest(cookie, SESSION_TOKEN)


@app.middleware("http")
async def auth_guard(request: Request, call_next):
    path = request.url.path
    open_paths = ("/api/login", "/healthz")
    if path.startswith("/api/") and path not in open_paths and not is_authed(request):
        return JSONResponse({"error": "unauthorized"}, status_code=401)
    return await call_next(request)


@app.get("/healthz")
def healthz():
    return {"status": "ok", "cache_configured": lang_cache is not None}


@app.post("/api/login")
def login(body: LoginBody, response: Response):
    if not hmac.compare_digest(body.password, APP_PASSWORD):
        return JSONResponse({"ok": False, "error": "invalid password"}, status_code=401)
    response.set_cookie(
        COOKIE_NAME, SESSION_TOKEN,
        httponly=True, samesite="lax", max_age=60 * 60 * 12, path="/",
    )
    return {"ok": True}


@app.get("/api/meta")
def meta():
    region = re.sub(r"^https?://", "", LANGCACHE_BASE_URL).split(".")[0]
    sdk_version = getattr(langcache_pkg, "__version__", None) if langcache_pkg else None
    if not sdk_version and langcache_pkg:
        try:
            from importlib.metadata import version
            sdk_version = version("langcache")
        except Exception:
            sdk_version = "unknown"
    cache_id_masked = (LANGCACHE_CACHE_ID[:6] + "…" + LANGCACHE_CACHE_ID[-4:]) if LANGCACHE_CACHE_ID else "—"
    return {
        "service": {
            "base_url": LANGCACHE_BASE_URL,
            "region": region,
            "cache_id_masked": cache_id_masked,
            "sdk_version": sdk_version or "n/a",
            "llm_model": OPENAI_MODEL,
            "cache_configured": lang_cache is not None,
        },
        "stats": stats_snapshot(),
    }


@app.post("/api/flush")
def flush(body: FlushBody):
    attrs = build_attributes(body.company.strip(), body.bu.strip(), body.person.strip(), body.isolation)
    if not attrs:
        return JSONResponse(
            {"ok": False, "error": "Global scope cannot be flushed by attributes. Pick a narrower scope."},
            status_code=400,
        )
    if not lang_cache:
        return JSONResponse({"ok": False, "error": "LangCache is not configured."}, status_code=503)
    try:
        res = lang_cache.delete_query(attributes=attrs)
        deleted = getattr(res, "deleted_entries_count", None)
        return {"ok": True, "attributes": attrs, "deleted": deleted}
    except Exception as e:
        return JSONResponse({"ok": False, "attributes": attrs, "error": str(e)}, status_code=500)


@app.post("/api/chat")
def chat(body: ChatBody):
    prompt = (body.prompt or "").strip()
    if not prompt:
        return JSONResponse({"error": "empty prompt"}, status_code=400)

    company = (body.company or "Acme").strip()
    bu = (body.bu or "BU-1").strip()
    person = (body.person or "user-1").strip()
    attrs = build_attributes(company, bu, person, body.isolation)
    ttl_ms = body.ttl_seconds * 1000 if body.ttl_seconds > 0 else None
    sim_threshold = None if body.similarity_threshold < 0 else body.similarity_threshold

    STATS["requests"] += 1
    trace: List[Dict[str, Any]] = []
    debug: Dict[str, Any] = {
        "attributes": attrs,
        "isolation": body.isolation,
        "similarity_threshold": sim_threshold,
        "strategy": "EXACT→SEMANTIC fallback" if body.exact_then_semantic else "SEMANTIC only",
        "ttl_ms_on_set": ttl_ms,
    }

    # --- Profile write: "My role is X" / "Minha função é X" ---
    role_set = try_extract_role_set(prompt)
    if role_set and lang_cache:
        t0 = time.perf_counter()
        try:
            lang_cache.set(prompt=KEY_ROLE, response=f"Your role is {role_set}.", attributes=attrs, ttl_millis=ttl_ms)
            set_ms = (time.perf_counter() - t0) * 1000
            STATS["stores"] += 1
            debug.update({"intent": "identity:role:set", "stored_under": KEY_ROLE, "value": role_set})
            trace.append({"label": "LangCache SET", "ms": set_ms, "status": "stored"})
            answer = f"Got it! Your role ({role_set}) was stored in the cache, for your current scope only."
            return {
                "answer": answer, "source": "stored", "intent": "identity:role:set",
                "cache_key": KEY_ROLE, "attributes": attrs,
                "cache_ms": set_ms, "llm_ms": None, "similarity": None, "entry_id": None,
                "strategy": None, "tokens_est": estimate_tokens(prompt, answer), "saved_usd": 0.0,
                "trace": trace, "debug": debug, "stats": stats_snapshot(),
            }
        except Exception as e:
            return JSONResponse({"error": f"Failed to store role: {e}", "debug": debug}, status_code=500)

    cache_key, intent = normalize_prompt_for_cache(prompt)
    debug["intent"] = intent

    # --- Name questions are answered directly and never cached ---
    if intent == "identity:name":
        answer = (
            f"Hi, {person}! Your name is {person}."
            if person else
            "I don't have your name saved. Fill in the Person field in the panel."
        )
        trace.append({"label": "Direct answer (the name never goes to the cache)", "ms": 0, "status": "direct"})
        return {
            "answer": answer, "source": "direct", "intent": intent,
            "cache_key": KEY_NAME, "attributes": attrs,
            "cache_ms": None, "llm_ms": None, "similarity": None, "entry_id": None,
            "strategy": None, "tokens_est": estimate_tokens(prompt, answer), "saved_usd": 0.0,
            "trace": trace, "debug": debug, "stats": stats_snapshot(),
        }

    # --- Role questions use EXACT search only ---
    strategies = None
    if intent == "identity:role":
        strategies = [SearchStrategy.EXACT] if SearchStrategy is not None else None
        sim_threshold = None

    # --- Ambiguity: rewrite the prompt with the persona's domain ---
    rewritten = prompt
    domain_label = infer_domain(company, bu)
    if looks_ambiguous(prompt):
        rewritten = rewrite_with_domain(prompt, domain_label)
        if intent == "fact":
            cache_key = f"[FACT]\n{rewritten}"
        debug["domain_rewrite"] = rewritten
    debug["cache_key"] = cache_key

    # --- 1) LangCache search ---
    hit = None
    cache_ms = None
    if lang_cache:
        t0 = time.perf_counter()
        try:
            if strategies is not None:
                results = lang_cache.search(prompt=cache_key, attributes=attrs, search_strategies=strategies)
            else:
                results = lang_cache.search(prompt=cache_key, similarity_threshold=sim_threshold, attributes=attrs)
            if (not results or not getattr(results, "data", None)) and body.exact_then_semantic \
                    and SearchStrategy is not None and strategies is None:
                results = lang_cache.search(
                    prompt=cache_key, attributes=attrs,
                    search_strategies=[SearchStrategy.EXACT, SearchStrategy.SEMANTIC],
                )
            if results and getattr(results, "data", None):
                hit = results.data[0]
        except Exception as e:
            debug["cache_search_error"] = str(e)
        cache_ms = (time.perf_counter() - t0) * 1000
        STATS["cache_ms"].append(cache_ms)
        trace.append({"label": "LangCache SEARCH", "ms": cache_ms, "status": "hit" if hit else "miss"})

    # --- 2) Cache hit: return immediately ---
    if hit is not None:
        answer = depersonalize_safe(hit.response, person)
        tokens_est = estimate_tokens(prompt, answer)
        saved = calc_savings(tokens_est, body.price_in_per_m, body.price_out_per_m, body.frac_in)
        STATS["hits"] += 1
        STATS["tokens_saved"] += tokens_est
        STATS["usd_saved"] += saved
        similarity = getattr(hit, "similarity", None)
        strategy_used = getattr(hit, "search_strategy", None) or getattr(hit, "searchStrategy", None)
        if strategy_used is not None:
            strategy_used = str(getattr(strategy_used, "value", strategy_used)).lower()
        return {
            "answer": answer, "source": "cache", "intent": intent,
            "cache_key": cache_key, "attributes": attrs,
            "cache_ms": cache_ms, "llm_ms": None,
            "similarity": similarity, "entry_id": getattr(hit, "id", None), "strategy": strategy_used,
            "tokens_est": tokens_est, "saved_usd": saved,
            "trace": trace, "debug": debug, "stats": stats_snapshot(),
        }

    # --- 3) Miss: ask the LLM (or explain missing role), then store ---
    llm_ms = None
    if intent == "identity:role":
        answer = "I don't have your role in this scope yet. Say \"My role is <title>\" and I will store it."
        trace.append({"label": "No role stored in this scope", "ms": 0, "status": "direct"})
        stored = False
    else:
        t1 = time.perf_counter()
        answer = call_openai(rewritten, company, bu)
        llm_ms = (time.perf_counter() - t1) * 1000
        STATS["llm_ms"].append(llm_ms)
        trace.append({"label": f"OpenAI {OPENAI_MODEL}", "ms": llm_ms, "status": "llm"})
        stored = False
        if lang_cache:
            t2 = time.perf_counter()
            try:
                lang_cache.set(prompt=cache_key, response=answer, attributes=attrs, ttl_millis=ttl_ms)
                stored = True
                STATS["stores"] += 1
                trace.append({"label": "LangCache SET", "ms": (time.perf_counter() - t2) * 1000, "status": "stored"})
            except Exception as e:
                debug["cache_set_error"] = str(e)

    STATS["misses"] += 1
    answer = depersonalize_safe(answer, person)
    return {
        "answer": answer, "source": "llm", "intent": intent,
        "cache_key": cache_key, "attributes": attrs,
        "cache_ms": cache_ms, "llm_ms": llm_ms,
        "similarity": None, "entry_id": None, "strategy": None,
        "stored": stored,
        "tokens_est": estimate_tokens(prompt, answer), "saved_usd": 0.0,
        "trace": trace, "debug": debug, "stats": stats_snapshot(),
    }


# Static UI (mounted last so /api/* wins)
@app.get("/")
def index():
    return FileResponse(os.path.join(STATIC_DIR, "index.html"))


app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")


if __name__ == "__main__":
    port = int(os.getenv("PORT", "7860"))
    uvicorn.run(app, host="0.0.0.0", port=port)