import tiktoken

def count_tokens(text: str, model: str = "gpt-4o") -> int:
    encoding = tiktoken.encoding_for_model(model)
    return len(encoding.encode(text))

# 检查 prompt 消耗
system_prompt = "你是一个客服助手，只回答关于产品使用的问题，不讨论其他话题。"
print(f"System prompt token 数：{count_tokens(system_prompt)}")

import os
import time
import tiktoken
from openai import OpenAI
from dataclasses import dataclass, field
from typing import List, Dict, Optional

client = OpenAI(base_url="https://api.deepseek.com", api_key=os.getenv("DEEPSEEK_API_KEY"))

# 各模型的定价（美元/百万 token，仅作参考，实际使用请查最新官网）
# ⚠️ 定价随时可能变化，使用前请核实：https://platform.deepseek.com/api-docs/pricing
MODEL_PRICING = {
    # DeepSeek（课程主要使用）
    "deepseek-chat": {"input": 0.27, "output": 1.10},      # DeepSeek-V3，截至2026年3月
    "deepseek-reasoner": {"input": 0.55, "output": 2.19},  # DeepSeek-R1，截至2026年3月
    # OpenAI（部分章节使用）
    "gpt-4o": {"input": 2.5, "output": 10.0},
    "gpt-4o-mini": {"input": 0.15, "output": 0.60},
}


@dataclass
class CostTracker:
    """成本追踪器，记录每次调用的 token 和费用"""
    total_input_tokens: int = 0
    total_output_tokens: int = 0
    total_cost_usd: float = 0.0
    call_count: int = 0
    call_history: List[Dict] = field(default_factory=list)

    def record(self, model: str, input_tokens: int, output_tokens: int):
        pricing = MODEL_PRICING.get(model, {"input": 2.5, "output": 10.0})
        cost = (input_tokens * pricing["input"] + output_tokens * pricing["output"]) / 1_000_000

        self.total_input_tokens += input_tokens
        self.total_output_tokens += output_tokens
        self.total_cost_usd += cost
        self.call_count += 1
        self.call_history.append({
            "model": model,
            "input_tokens": input_tokens,
            "output_tokens": output_tokens,
            "cost_usd": cost,
            "timestamp": time.time()
        })

    def summary(self) -> str:
        return (
            f"调用次数：{self.call_count}\n"
            f"总 input token：{self.total_input_tokens:,}\n"
            f"总 output token：{self.total_output_tokens:,}\n"
            f"总费用：${self.total_cost_usd:.4f} USD（约 ¥{self.total_cost_usd * 7.2:.2f}）"
        )


# 全局追踪器
tracker = CostTracker()


def chat_with_history(
    user_message: str,
    history: List[Dict],
    system_prompt: str,
    model: str = "deepseek-chat",
    max_history_turns: int = 10,
    max_output_tokens: int = 500,
    rag_context: Optional[str] = None,
) -> tuple[str, List[Dict]]:
    """
    带成本控制的多轮对话函数

    优化点：
    1. 对话历史截断（max_history_turns）
    2. 输出长度限制（max_output_tokens）
    3. RAG 内容可选注入
    4. 自动成本追踪
    """

    # 截断对话历史，只保留最近 N 轮（每轮 = 用户 + 助手，共 2 条消息）
    truncated_history = history[-(max_history_turns * 2):]
    if len(history) > len(truncated_history):
        print(f"[成本控制] 历史截断：{len(history)} 条 → {len(truncated_history)} 条")

    # 构建 system prompt（如果有 RAG 内容，追加到 system prompt 末尾）
    full_system = system_prompt
    if rag_context:
        # RAG 内容截断到 1500 token
        encoding = tiktoken.encoding_for_model(model if model in tiktoken.list_encoding_names() else "gpt-4o")
        rag_tokens = encoding.encode(rag_context)
        if len(rag_tokens) > 1500:
            rag_context = encoding.decode(rag_tokens[:1500])
            print(f"[成本控制] RAG 内容截断至 1500 token")
        full_system += f"\n\n参考资料：\n{rag_context}"

    messages = [{"role": "system", "content": full_system}]
    messages.extend(truncated_history)
    messages.append({"role": "user", "content": user_message})

    # 发送请求前预估 token（可选，用于监控）
    encoding = tiktoken.get_encoding("cl100k_base")
    estimated_input = sum(len(encoding.encode(m["content"])) for m in messages)
    print(f"[成本预估] 预计 input token：{estimated_input}")

    # 调用 API
    response = client.chat.completions.create(
        model=model,
        messages=messages,
        max_tokens=max_output_tokens,
        temperature=0.7,
    )

    assistant_message = response.choices[0].message.content
    usage = response.usage

    # 记录实际费用
    tracker.record(model, usage.prompt_tokens, usage.completion_tokens)
    print(
        f"[成本记录] 实际 input: {usage.prompt_tokens}, output: {usage.completion_tokens}, "
        f"本次费用: ${tracker.call_history[-1]['cost_usd']:.6f}"
    )

    # 更新对话历史
    new_history = truncated_history + [
        {"role": "user", "content": user_message},
        {"role": "assistant", "content": assistant_message},
    ]

    return assistant_message, new_history


# 使用示例
if __name__ == "__main__":
    system = "你是一个产品客服，只回答关于产品使用的问题，回答简洁不超过100字。"
    history = []

    questions = [
        "这个产品怎么退款？",
        "退款要多久到账？",
        "如果超过7天还没到账怎么办？",
    ]

    for q in questions:
        print(f"\n用户：{q}")
        answer, history = chat_with_history(
            user_message=q,
            history=history,
            system_prompt=system,
            model="deepseek-chat",
            max_history_turns=5,
            max_output_tokens=200,
        )
        print(f"助手：{answer}")

    print(f"\n{'='*40}")
    print("成本汇总：")
    print(tracker.summary())

import os
os.environ["LANGCHAIN_TRACING_V2"] = "true"
os.environ["LANGCHAIN_API_KEY"] = "your-langsmith-api-key"
# 之后所有通过 LangChain 的调用都会自动上报