File size: 1,347 Bytes
5dbe83d 752e673 5dbe83d 09ee734 5dbe83d 65b412a 5dbe83d 752e673 65b412a 5dbe83d 65b412a 5dbe83d 65b412a 5dbe83d 65b412a 5dbe83d 5aea999 5dbe83d 65b412a 5dbe83d 65b412a 09ee734 8d942c4 65b412a 8d942c4 65b412a 09ee734 65b412a 8d942c4 5dbe83d 65b412a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 |
import os
import pathlib
from typing import NewType
import pydantic
import player
import message
from pymongo import MongoClient
from pydantic import BaseModel
DATA_COLLECTION_MODE = os.environ.get("DATA_COLLECTION_MODE", "JSONL")
MONGODB_CONNECTION_STRING = os.environ.get("MONGODB_CONNECTION_STRING")
DB_NAME = os.environ.get("MONGODB_NAME")
Model = NewType("Model", BaseModel)
def save(log_object: Model):
collection = get_collection(log_object)
if DATA_COLLECTION_MODE.upper() == "JSONL":
data_dir = pathlib.Path(__file__).parent.parent / "data"
log_file = os.path.join(data_dir, f"{collection}.jsonl")
with open(log_file, "a+") as f:
f.write(log_object.model_dump_json() + "\n")
if DATA_COLLECTION_MODE.upper() == "MONGODB":
client = MongoClient(MONGODB_CONNECTION_STRING)
db = client[DB_NAME]
db[collection].insert_one(log_object.model_dump())
def get_collection(log_object: Model) -> str:
from game import Game
if isinstance(log_object, message.AgentMessage):
collection = "messages"
elif isinstance(log_object, player.Player):
collection = "players"
elif isinstance(log_object, Game):
collection = "games"
else:
raise ValueError(f"Unknown log object type: {type(log_object)}")
return collection
|