Cost Optimization

"""Đếm token và ước lượng chi phí trước khi gọi LLM API.

Chức năng:
  - Đếm token chính xác bằng tokenizer chính thức (tiktoken cho OpenAI).
  - Ước lượng chi phí USD theo bảng giá cập nhật.
  - Từ chối request nếu vượt budget/request hoặc budget/ngày.
  - Ghi log để dashboard chi phí phân tích sau.
"""

from __future__ import annotations

import json
import time
from dataclasses import dataclass, asdict
from typing import Iterable

import tiktoken  # pip install tiktoken


# Bảng giá USD / 1K token — đồng bộ với bảng chính thức trước khi dùng prod.
MODEL_PRICES = {
    "gpt-4o":          {"in": 0.005,    "out": 0.015},
    "gpt-4o-mini":     {"in": 0.00015,  "out": 0.00060},
    "claude-3-5-sonnet": {"in": 0.003,  "out": 0.015},
    "claude-3-haiku":  {"in": 0.00025,  "out": 0.00125},
    "gemini-1.5-flash": {"in": 0.000075,"out": 0.0003},
}


@dataclass
class TokenEstimate:
    model: str
    input_tokens: int
    estimated_output_tokens: int
    cost_input_usd: float
    cost_output_usd: float
    total_usd: float

    def to_json(self) -> str:
        return json.dumps(asdict(self), ensure_ascii=False)


def count_tokens(text: str, model: str = "gpt-4o") -> int:
    """Đếm token chính xác bằng tokenizer của model tương ứng."""
    try:
        enc = tiktoken.encoding_for_model(model)
    except KeyError:
        enc = tiktoken.get_encoding("cl100k_base")
    return len(enc.encode(text))


def estimate_cost(
    prompt: str,
    model: str = "gpt-4o",
    expected_output_tokens: int = 300,
) -> TokenEstimate:
    """Ước lượng chi phí dựa trên prompt và số output dự kiến."""
    if model not in MODEL_PRICES:
        raise ValueError(f"Unknown model: {model}")

    price = MODEL_PRICES[model]
    input_tokens = count_tokens(prompt, model)

    cost_in = input_tokens / 1000 * price["in"]
    cost_out = expected_output_tokens / 1000 * price["out"]
    total = cost_in + cost_out

    return TokenEstimate(
        model=model,
        input_tokens=input_tokens,
        estimated_output_tokens=expected_output_tokens,
        cost_input_usd=round(cost_in, 6),
        cost_output_usd=round(cost_out, 6),
        total_usd=round(total, 6),
    )


class BudgetGuard:
    """Từ chối request nếu vượt budget/request hoặc budget/ngày."""

    def __init__(
        self,
        max_per_request_usd: float = 0.50,
        max_per_day_usd: float = 500.0,
    ) -> None:
        self.max_per_request = max_per_request_usd
        self.max_per_day = max_per_day_usd
        self._spent_today_usd = 0.0
        self._day_start = time.time()

    def _reset_if_new_day(self) -> None:
        if time.time() - self._day_start > 86400:
            self._spent_today_usd = 0.0
            self._day_start = time.time()

    def check(self, estimate: TokenEstimate) -> None:
        self._reset_if_new_day()
        if estimate.total_usd > self.max_per_request:
            raise RuntimeError(
                f"Request cost $"
                f"{estimate.total_usd:.4f} > limit $"
                f"{self.max_per_request:.2f}"
            )
        if self._spent_today_usd + estimate.total_usd > self.max_per_day:
            raise RuntimeError(
                f"Daily budget exceeded: $"
                f"{self._spent_today_usd + estimate.total_usd:.2f} > $"
                f"{self.max_per_day:.2f}"
            )

    def record(self, estimate: TokenEstimate) -> None:
        self._spent_today_usd += estimate.total_usd


# --- Ví dụ sử dụng ---
if __name__ == "__main__":
    prompt = "Tóm tắt báo cáo Q3 này thành 5 gạch đầu dòng: ..."
    est = estimate_cost(prompt, model="gpt-4o", expected_output_tokens=200)
    print(est.to_json())

    guard = BudgetGuard(max_per_request_usd=0.10, max_per_day_usd=100.0)
    try:
        guard.check(est)
        # ... gọi LLM thật ở đây ...
        guard.record(est)
    except RuntimeError as e:
        print("Rejected:", e)

"""Pipeline tối ưu chi phí cho chatbot:
    1) Semantic cache (Redis + embedding).
    2) Classifier định tuyến small vs flagship model.
    3) Ghi log chi phí để dashboard phân tích.
"""

from __future__ import annotations

import hashlib
import json
import os
import time
from typing import Literal, Optional

import numpy as np
import redis
from sentence_transformers import SentenceTransformer


EMBEDDER = SentenceTransformer("all-MiniLM-L6-v2")
CACHE = redis.Redis.from_url(os.environ.get("REDIS_URL", "redis://localhost:6379"))
SIMILARITY_THRESHOLD = 0.95


def _embed(text: str) -> np.ndarray:
    return EMBEDDER.encode(text, normalize_embeddings=True)


def _cache_key(prefix: str, text: str) -> str:
    digest = hashlib.sha1(text.encode("utf-8")).hexdigest()[:16]
    return f"{prefix}:{digest}"


def semantic_cache_lookup(query: str) -> Optional[str]:
    """Tìm câu trả lời gần nhất trong cache. Trả None nếu miss."""
    query_emb = _embed(query).astype(np.float32)

    # Quét tối đa 200 entry gần nhất — production nên dùng FAISS / pgvector.
    for key in CACHE.scan_iter(match="cache:*", count=200):
        raw = CACHE.get(key)
        if raw is None:
            continue
        entry = json.loads(raw)
        cached_emb = np.asarray(entry["embedding"], dtype=np.float32)
        similarity = float(np.dot(query_emb, cached_emb))
        if similarity >= SIMILARITY_THRESHOLD:
            return entry["response"]
    return None


def semantic_cache_store(query: str, response: str, ttl: int = 3600) -> None:
    key = _cache_key("cache", query)
    payload = {
        "embedding": _embed(query).astype(np.float32).tolist(),
        "response": response,
        "stored_at": time.time(),
    }
    CACHE.setex(key, ttl, json.dumps(payload))


# ---------------------------------------------------------------------------
# Classifier định tuyến — trong thực tế có thể là logistic regression, small
# LM, hoặc gọi một API rẻ (Haiku). Ở đây mô phỏng bằng heuristic độ dài.
# ---------------------------------------------------------------------------

def classify_complexity(query: str) -> Literal["simple", "complex"]:
    word_count = len(query.split())
    has_code_hint = any(k in query.lower() for k in ["code", "python", "sql"])
    if word_count < 40 and not has_code_hint:
        return "simple"
    return "complex"


def call_small_model(query: str) -> str:
    # ... gọi Haiku / Flash / Llama-8B self-host ...
    return f"[small] Trả lời cho: {query[:40]}..."


def call_flagship_model(query: str) -> str:
    # ... gọi GPT-4o / Claude Sonnet ...
    return f"[flagship] Trả lời cho: {query[:40]}..."


def chat(query: str) -> dict:
    """Pipeline chính."""
    # 1. Thử cache.
    cached = semantic_cache_lookup(query)
    if cached:
        return {"response": cached, "source": "cache", "model": None}

    # 2. Route theo độ phức tạp.
    complexity = classify_complexity(query)
    if complexity == "simple":
        answer = call_small_model(query)
        model = "claude-haiku"
    else:
        answer = call_flagship_model(query)
        model = "gpt-4o"

    # 3. Lưu cache để lần sau khỏi gọi nữa.
    semantic_cache_store(query, answer)
    return {"response": answer, "source": "llm", "model": model}


if __name__ == "__main__":
    for q in [
        "Hà Nội hôm nay thời tiết thế nào?",
        "Thời tiết ở Hà Nội hôm nay ra sao?",
        "Viết Python function tính fibonacci bằng DP bottom-up",
    ]:
        r = chat(q)
        print(q, "->", r["source"], r["model"])