import wandb
import weave
import random

# 同じ project で W&B と Weave の両方を初期化します
project = "my-workspace-project"
weave.init(project)

@weave.op
def evaluate_model(model_state, epoch):
    # 評価メトリクスをシミュレートします
    accuracy = 0.7 + epoch * 0.02 + random.uniform(-0.05, 0.05)
    
    # モデルの特定の動作をトラッキングします
    test_responses = {
        "conciseness_check": "Model generated 500 words for simple question",
        "accuracy_check": "Model answered 8/10 questions correctly",
        "hallucination_check": "No factual errors detected"
    }
    
    return {
        "accuracy": accuracy,
        "diagnostic_results": test_responses
    }

@weave.op
def training_step(epoch, lr):
    # トレーニングのロジックをシミュレートします
    loss = 1.0 / (epoch + 1) + random.uniform(-0.1, 0.1)
    
    # トレースを伴う評価を行います
    eval_results = evaluate_model(f"checkpoint_{epoch}", epoch)
    
    return {
        "loss": loss,
        "accuracy": eval_results["accuracy"],
        "diagnostics": eval_results["diagnostic_results"]
    }

# トレーニングループ
with wandb.init(project=project, config={"lr": 0.01}) as run:
    for epoch in range(5):
        # Weave のトレースを使用してトレーニングを実行します
        results = training_step(epoch, run.config.lr)
        
        # W&B にログします - インテグレーションの接点を作成します
        run.log({
            "epoch": epoch,
            "loss": results["loss"],
            "accuracy": results["accuracy"]
        })

client.clear_wandb_run_context()

my_function(20)  # このトレースは、グローバルな wandb.run があれば使用します。