import os
from google import genai

# 使用个人 API Key 初始化
client = genai.Client(api_key=os.environ["GOOGLE_API_KEY"])

from google import genai

# 使用 Vertex AI 企业后端初始化
client = genai.Client(
    vertexai=True,
    project="my-gcp-enterprise-project",
    location="us-central1",
)

# 统一 API 调用 Claude 3.5 Sonnet
claude_response = client.models.generate_content(
    model="claude-3-5-sonnet-v2@20241022",
    contents="用简明的技术语言解释一下分布式事务的 2PC 机制。",
)
print(claude_response.text)

# 统一 API 调用开源 Llama 3.1
llama_response = client.models.generate_content(
    model="meta/llama-3.1-405b-instruct-maas",
    contents="Explain raft consensus algorithm briefly.",
)
print(llama_response.text)

import os
from google import genai
from google.genai import types

client = genai.Client(api_key=os.environ["GOOGLE_API_KEY"])

response = client.models.generate_content(
    model="gemini-2.0-flash",
    contents="def calculate_tax(amount): return amount * 0.1 有什么潜在隐患？",
    config=types.GenerateContentConfig(
        system_instruction="你是一个金融级支付系统架构师，专注于浮点数精度与边界异常审查。",
        temperature=0.1,
        max_output_tokens=1024,
    ),
)
print(response.text)

chat = client.chats.create(model="gemini-2.0-flash")

# 第一轮
r1 = chat.send_message("我们目前系统有 3 个微服务节点，日订单量约 50 万单。")
# 第二轮（模型自动携带历史上下文）
r2 = chat.send_message("基于刚才的体量，数据库读写分离是否是当务之急？")

print(r2.text)

from google import genai
from google.genai import types


def query_db_node_status(node_id: str) -> dict:
    """查询指定数据库节点的运行状态"""
    db = {
        "db-node-primary": {"role": "master", "status": "UP", "lag_ms": 0},
        "db-node-replica-1": {"role": "slave", "status": "UP", "lag_ms": 45},
    }
    return db.get(node_id, {"error": "节点不存在"})


# 1. 显式构建工具定义
tools_declaration = types.Tool(
    function_declarations=[
        types.FunctionDeclaration(
            name="query_db_node_status",
            description="查询数据库集群中节点的角色、健康状态与复制延迟",
            parameters=types.Schema(
                type=types.Type.OBJECT,
                properties={
                    "node_id": types.Schema(
                        type=types.Type.STRING,
                        description="节点 ID，如 'db-node-primary'",
                    )
                },
                required=["node_id"],
            ),
        )
    ]
)

tools_map = {"query_db_node_status": query_db_node_status}

# 2. 手动驱动状态机
contents = [
    types.Content(
        role="user",
        parts=[types.Part(text="检查 db-node-replica-1 节点的复制延迟情况。")],
    )
]

while True:
    response = client.models.generate_content(
        model="gemini-2.0-flash",
        contents=contents,
        config=types.GenerateContentConfig(tools=[tools_declaration]),
    )
    
    first_part = response.candidates[0].content.parts[0]
    
    # 检查是否触发了工具调用
    if hasattr(first_part, "function_call") and first_part.function_call:
        fc = first_part.function_call
        print(f"[拦截并执行工具] 函数: {fc.name}, 参数: {dict(fc.args)}")
        
        # 本地真实执行
        fn_output = tools_map[fc.name](**dict(fc.args))
        
        # 将模型的调用意图作为 model 消息追加
        contents.append(
            types.Content(
                role="model",
                parts=[types.Part(function_call=fc)],
            )
        )
        
        # 将执行结果作为 user 消息追加回上下文
        contents.append(
            types.Content(
                role="user",
                parts=[
                    types.Part(
                        function_response=types.FunctionResponse(
                            name=fc.name,
                            response={"result": fn_output},
                        )
                    )
                ],
            )
        )
    else:
        # 模型已输出最终文本
        print(f"\n[最终结论]\n{response.text}")
        break

def query_db_node_status(node_id: str) -> str:
    """查询指定数据库节点的运行状态与复制延迟。"""
    return f"节点 {node_id}: 状态正常, 复制延迟 12ms"


# 直接传函数，SDK 自动完成 Schema 提取与内部轮转
response = client.models.generate_content(
    model="gemini-2.0-flash",
    contents="帮我看看 db-node-replica-1 是否正常？",
    config=types.GenerateContentConfig(
        tools=[query_db_node_status],
        automatic_function_calling=types.AutomaticFunctionCallingConfig(disable=False),
    ),
)

print(response.text)

from google.genai import types

# 读取本地图片或架构图
with open("architecture_diagram.png", "rb") as f:
    img_data = f.read()

response = client.models.generate_content(
    model="gemini-2.0-flash",
    contents=[
        types.Part(
            inline_data=types.Blob(
                mime_type="image/png",
                data=img_data,
            )
        ),
        types.Part(text="请分析图中展现的架构拓扑，指出其单点故障（SPOF）风险。"),
    ],
)
print(response.text)