src/data_collection.py

import os
import pathlib
from typing import NewType

import pydantic

import message
import player

from pydantic import BaseModel

COLLECT_DATA = os.environ.get("COLLECT_DATA", "true").upper() == "TRUE"

Model = NewType("Model", BaseModel)


data_dir = pathlib.Path(__file__).parent.parent / "data"


def save(log_object: Model):
    if COLLECT_DATA:
        log_file = get_log_file(log_object)

        with open(log_file, "a+") as f:
            f.write(log_object.model_dump_json() + "\n")


def get_log_file(log_object: Model) -> str:
    if isinstance(log_object, message.AgentMessage):
        log_file = "messages.jsonl"
    elif isinstance(log_object, player.Player):
        log_file = "players.jsonl"
    else:
        raise ValueError(f"Unknown log object type: {type(log_object)}")

    return os.path.join(data_dir, log_file)