test_local.py

"""
无需 API Key 的本地测试
验证代码逻辑和数据结构
"""

import tiktoken
from pydantic import BaseModel, Field
from typing import Literal, List
import json


# ========== Token 计数测试 ==========

def test_token_counting():
    """测试 Token 计数功能"""
    print("=" * 60)
    print("Test 1: Token Counting")
    print("=" * 60)

    encoding = tiktoken.encoding_for_model("gpt-4")

    test_cases = [
        "Hello, how are you?",
        "你好，世界！",
        "The quick brown fox jumps over the lazy dog.",
        "人工智能正在改变世界",
    ]

    for text in test_cases:
        tokens = encoding.encode(text)
        print(f"\nText: {text}")
        print(f"Tokens: {tokens}")
        print(f"Count: {len(tokens)} tokens")

    print("\n[OK] Token counting works!")


# ========== Pydantic 模型测试 ==========

class OrderIntent(BaseModel):
    """订单意图"""
    order_id: str = Field(description="订单号")
    intent: Literal["query", "cancel", "modify"] = Field(description="意图")
    urgency: Literal["high", "medium", "low"] = Field(description="紧急程度")
    reason: str = Field(description="原因")


def test_pydantic_models():
    """测试 Pydantic 数据模型"""
    print("\n" + "=" * 60)
    print("Test 2: Pydantic Data Models")
    print("=" * 60)

    # 创建实例
    order = OrderIntent(
        order_id="ORD123456",
        intent="cancel",
        urgency="high",
        reason="Address is incorrect"
    )

    print("\nOrder Intent Object:")
    print(order)

    print("\nAs JSON:")
    print(order.model_dump_json(indent=2))

    print("\nAs Dict:")
    print(order.model_dump())

    # 验证数据校验
    print("\n--- Data Validation Test ---")
    try:
        invalid_order = OrderIntent(
            order_id="ORD999",
            intent="invalid_intent",  # 这会失败
            urgency="high",
            reason="Test"
        )
    except Exception as e:
        print(f"[Expected Error] Validation failed: {type(e).__name__}")

    print("\n[OK] Pydantic models work!")


# ========== 成本计算测试 ==========

PRICING = {
    "gpt-4": {"input": 0.03 / 1000, "output": 0.06 / 1000},
    "gpt-3.5-turbo": {"input": 0.0005 / 1000, "output": 0.0015 / 1000}
}


def calculate_cost(model: str, input_tokens: int, output_tokens: int) -> dict:
    """计算成本"""
    pricing = PRICING[model]
    input_cost = input_tokens * pricing["input"]
    output_cost = output_tokens * pricing["output"]
    total_cost = input_cost + output_cost

    return {
        "input_tokens": input_tokens,
        "output_tokens": output_tokens,
        "total_tokens": input_tokens + output_tokens,
        "cost_usd": total_cost,
        "cost_cny": total_cost * 7.2
    }


def test_cost_calculation():
    """测试成本计算"""
    print("\n" + "=" * 60)
    print("Test 3: Cost Calculation")
    print("=" * 60)

    test_cases = [
        ("gpt-3.5-turbo", 1000, 500),
        ("gpt-4", 1000, 500),
        ("gpt-4", 10000, 2000),
    ]

    print(f"\n{'Model':<20} {'Input':<10} {'Output':<10} {'Cost (USD)':<15} {'Cost (CNY)'}")
    print("-" * 75)

    for model, input_t, output_t in test_cases:
        cost = calculate_cost(model, input_t, output_t)
        print(f"{model:<20} {input_t:<10} {output_t:<10} "
              f"${cost['cost_usd']:<14.4f} "
              f"CNY {cost['cost_cny']:<.2f}")

    print("\n[OK] Cost calculation works!")


# ========== 消息结构测试 ==========

def test_message_structure():
    """测试 OpenAI 消息格式"""
    print("\n" + "=" * 60)
    print("Test 4: Message Structure")
    print("=" * 60)

    messages = [
        {
            "role": "system",
            "content": "You are a helpful AI assistant."
        },
        {
            "role": "user",
            "content": "What is machine learning?"
        },
        {
            "role": "assistant",
            "content": "Machine learning is a subset of AI..."
        },
        {
            "role": "user",
            "content": "Can you explain more?"
        }
    ]

    print("\nValid message format:")
    print(json.dumps(messages, indent=2, ensure_ascii=False))

    # 计算总 token 数
    encoding = tiktoken.encoding_for_model("gpt-4")
    total_tokens = 0
    for msg in messages:
        total_tokens += len(encoding.encode(msg["content"]))
        total_tokens += 4  # 消息格式开销

    print(f"\nEstimated tokens: {total_tokens}")
    print("\n[OK] Message structure works!")


# ========== 主函数 ==========

def main():
    """运行所有测试"""
    print("\n" + "=" * 60)
    print("Python Agent Tutorial - Local Tests")
    print("No API Key Required")
    print("=" * 60)

    try:
        test_token_counting()
        test_pydantic_models()
        test_cost_calculation()
        test_message_structure()

        print("\n" + "=" * 60)
        print("[SUCCESS] All local tests passed!")
        print("=" * 60)
        print("\nNext steps:")
        print("1. Set up your OPENAI_API_KEY in .env file")
        print("2. Run: python examples/01-basic-chat/simple_chat.py")
        print("3. Continue with the tutorials")

    except Exception as e:
        print(f"\n[ERROR] Test failed: {e}")
        import traceback
        traceback.print_exc()


if __name__ == "__main__":
    main()