import os
import anthropic

client = anthropic.Anthropic(api_key=os.environ["ANTHROPIC_API_KEY"])

response = client.messages.create(
    model="claude-3-5-sonnet-20241022",
    max_tokens=1024,
    system="你是一名资深的分布式高可用审计员，专注于网络分区与脑裂隐患。",
    messages=[
        {"role": "user", "content": "etcd 在遭遇网络分区时，少数派节点集群如何保证强一致性？"}
    ],
)

# content 是列表，提取首个文本块
print(response.content[0].text)

import json
import anthropic

client = anthropic.Anthropic()

# 1. 声明工具的 JSON Schema
tools_schema = [
    {
        "name": "query_slow_queries",
        "description": "查询指定数据库实例中超过阈值的慢查询日志列表",
        "input_schema": {
            "type": "object",
            "properties": {
                "instance_id": {"type": "string", "description": "数据库实例标识"},
                "duration_threshold_ms": {"type": "integer", "description": "耗时阈值，毫秒"},
            },
            "required": ["instance_id", "duration_threshold_ms"],
        },
    }
]


# 2. 本地真实函数
def query_slow_queries(instance_id: str, duration_threshold_ms: int) -> str:
    # 模拟日志系统返回
    return json.dumps([
        {
            "query": "SELECT * FROM orders WHERE user_id = ? AND status = ?",
            "cost_ms": 2450,
            "rows_scanned": 1500000,
            "index_used": None,
        }
    ])


tools_registry = {"query_slow_queries": query_slow_queries}


# 3. 显式驱动循环
def run_database_diagnosis(user_prompt: str) -> str:
    messages = [{"role": "user", "content": user_prompt}]
    
    while True:
        response = client.messages.create(
            model="claude-3-5-sonnet-20241022",
            max_tokens=2048,
            tools=tools_schema,
            messages=messages,
        )
        
        # 判断模型的停止原因
        if response.stop_reason == "tool_use":
            # 1. 提取所有工具调用块（支持单轮内模型并行调用多个工具）
            tool_use_blocks = [b for b in response.content if b.type == "tool_use"]
            tool_results = []
            
            for tub in tool_use_blocks:
                print(f"[审计拦截] 模型请求执行工具: {tub.name}, 参数: {tub.input}")
                
                # 执行本地逻辑
                output_str = tools_registry[tub.name](**tub.input)
                
                # 组装针对该 tool_use_id 的执行产物
                tool_results.append({
                    "type": "tool_result",
                    "tool_use_id": tub.id,  # 必须与调用块中的 ID 精确对应
                    "content": output_str,
                })
            
            # 2. 维持消息上下文的严格对称性
            # 首先追加模型给出的 assistant 响应（包含 tool_use 声明）
            messages.append({"role": "assistant", "content": response.content})
            # 紧接着追加包含所有执行结果的 user 消息
            messages.append({"role": "user", "content": tool_results})
            
        elif response.stop_reason == "end_turn":
            # 模型已收敛并给出最终结论
            for block in response.content:
                if hasattr(block, "text"):
                    return block.text
            return "无文本产出"
        else:
            raise RuntimeError(f"异常停止状态: {response.stop_reason}")


# 执行调用
conclusion = run_database_diagnosis(
    "请诊断 db-orders-01 实例，找出耗时超过 2000ms 的慢查询并给出优化方案。"
)
print(f"\n================ 诊断建议 ================\n{conclusion}")

response = client.messages.create(
    model="claude-3-7-sonnet-20250219",
    max_tokens=16000,
    thinking={
        "type": "enabled",
        "budget_tokens": 4000,  # 允许模型思考最多消耗 4000 tokens
    },
    messages=[{
        "role": "user",
        "content": "我们需要为千万级日活的交易网关设计限流降级方案，对比令牌桶、滑动窗口与分布式 Redis 方案的利弊并给出具体架构。",
    }],
)

# 响应解析：思考内容与最终回答天然隔离
for block in response.content:
    if block.type == "thinking":
        print(f"--- 内部推演过程（前 300 字）---\n{block.thinking[:300]}...\n")
    elif block.type == "text":
        print(f"--- 最终架构报告 ---\n{block.text}")