ryanbalch's picture
finalizing top level metadata; storing things in vectorstore
4c91492
import os
import json
from pydantic import BaseModel, Field
from typing import Literal, Optional, ClassVar
from slugify import slugify
class Event(BaseModel):
id: str = Field(default_factory=slugify)
game_name: str
description: str
event: str
minute: str
player: str
team: str
@classmethod
def from_row(cls, game_name: str, row_id: int, **kwargs):
return cls(
id=f"{game_name}_{row_id}",
game_name=game_name,
description=kwargs["description"],
event=kwargs["event"],
minute=kwargs["minute"],
player=kwargs["player"],
team=kwargs["team"],
)
def event_vector_metadata(self):
return {
"type": "event",
"game_name": self.game_name,
"description": self.description,
"event": self.event,
"minute": self.minute,
"player": self.player,
"team": self.team,
}
@classmethod
def get_events(cls):
for filename in os.listdir("/workspace/data/huge-league/games"):
with open(os.path.join("/workspace/data/huge-league/games", filename), 'r') as f:
data = json.load(f)
for row_id, event in enumerate(data):
yield cls.from_row(game_name=filename.split(".")[0], row_id=row_id, **event)