mirror of
https://github.com/velocitatem/PHANTOM.git
synced 2026-05-31 08:33:36 +00:00
115 lines
3.7 KiB
Python
115 lines
3.7 KiB
Python
import os
|
|
from pydantic import BaseModel as Base
|
|
import json
|
|
|
|
class PayloadModel(Base):
|
|
sessionId: str
|
|
experimentId: str | None
|
|
eventName: str
|
|
page: str | None
|
|
productId: str | None
|
|
metadata: dict
|
|
storeMode: str
|
|
userAgent: str
|
|
ts: str
|
|
|
|
class ValueModel(Base):
|
|
payload: PayloadModel
|
|
encoding: str
|
|
isPayloadNull: bool
|
|
schemaId: int
|
|
size: int
|
|
|
|
class InteractionModel(Base):
|
|
partitionID: int
|
|
offset: int
|
|
timestamp: int
|
|
compression: str
|
|
isTransactional: bool
|
|
headers: list
|
|
key: dict
|
|
value: ValueModel
|
|
|
|
class Loader:
|
|
def __init__(self, src_dir: str):
|
|
self.src_dir = src_dir
|
|
self.entries = os.listdir(src_dir)
|
|
if not self.entries: raise ValueError("empty directory")
|
|
self.data = self._load_sessions()
|
|
|
|
def _is_admin_page(self, interaction: InteractionModel) -> bool:
|
|
page = interaction.value.payload.page
|
|
return page and page.startswith("/admin/")
|
|
|
|
def _load_sessions(self) -> dict:
|
|
sessions = {}
|
|
for entry in self.entries:
|
|
int_path = f"{self.src_dir}/{entry}/int.json"
|
|
raw = json.load(open(int_path))
|
|
ints = [InteractionModel(**i) for i in raw]
|
|
sessions[entry] = [i for i in ints if not self._is_admin_page(i)]
|
|
return sessions
|
|
|
|
def get_data(self) -> dict:
|
|
return self.data
|
|
|
|
def get_entries(self) -> tuple[list[str], int]:
|
|
return self.entries, len(self.entries)
|
|
|
|
class AgentLoader(Loader):
|
|
"""Loader for agent interaction data with simplified schema (direct PayloadModel format)"""
|
|
|
|
def _is_admin_page_simple(self, interaction: PayloadModel) -> bool:
|
|
return interaction.page and interaction.page.startswith("/admin/")
|
|
|
|
def _load_sessions(self) -> dict:
|
|
sessions = {}
|
|
for entry in self.entries:
|
|
int_path = f"{self.src_dir}/{entry}/int.json"
|
|
raw = json.load(open(int_path))
|
|
ints = [PayloadModel(**i) for i in raw]
|
|
sessions[entry] = [i for i in ints if not self._is_admin_page_simple(i)]
|
|
return sessions
|
|
|
|
class JointLoader:
|
|
"""Loader for combined human (Kafka) and agent (direct) data without discrimination"""
|
|
|
|
def __init__(self, human_dir: str, agent_dir: str):
|
|
self.human_dir = human_dir
|
|
self.agent_dir = agent_dir
|
|
self.human_loader = Loader(human_dir)
|
|
self.agent_loader = AgentLoader(agent_dir)
|
|
self.data = self._load_joint_sessions()
|
|
self.entries = list(self.data.keys())
|
|
|
|
def _load_joint_sessions(self) -> dict:
|
|
sessions = {}
|
|
# load human sessions (unwrap from Kafka format to PayloadModel)
|
|
for sid, evts in self.human_loader.get_data().items():
|
|
sessions[f"human_{sid}"] = [evt.value.payload for evt in evts]
|
|
# load agent sessions (already PayloadModel)
|
|
for sid, evts in self.agent_loader.get_data().items():
|
|
sessions[f"agent_{sid}"] = evts
|
|
return sessions
|
|
|
|
def get_data(self) -> dict:
|
|
return self.data
|
|
|
|
def get_entries(self) -> tuple[list[str], int]:
|
|
return self.entries, len(self.entries)
|
|
|
|
if __name__ == "__main__":
|
|
AGENT_DIR = "/home/velocitatem/Documents/Projects/PHANTOM/experiments/agents/collected_data/"
|
|
loader = AgentLoader(AGENT_DIR)
|
|
_, n = loader.get_entries()
|
|
print(f"Loaded {n} agent sessions from {AGENT_DIR}")
|
|
|
|
HUMAN_DIR = "/home/velocitatem/Documents/Projects/PHANTOM/experiments/collected_data/"
|
|
loader = Loader(HUMAN_DIR)
|
|
_, n = loader.get_entries()
|
|
print(f"Loaded {n} human sessions from {HUMAN_DIR}")
|
|
|
|
joint_loader = JointLoader(HUMAN_DIR, AGENT_DIR)
|
|
_, n = joint_loader.get_entries()
|
|
print(f"Loaded {n} total sessions (combined) from joint loader")
|