Spaces:
Sleeping
Sleeping
File size: 4,542 Bytes
b88fdd0 | 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 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 | import pandas as pd
import asyncio
import json
import os
import time
from typing import Optional, Generator, AsyncGenerator
from pydantic import BaseModel
class SimulationConfig(BaseModel):
speed: float = 1.0 # Multiplier, or seconds delay? Let's say speed multiplier (1x, 2x...)
# But for simplicity, let's treat it as "transactions per second" or delay.
# The requirement says "Speed Slider: Adjusts how fast...".
# Let's map speed 1-10 to delay.
# Speed 1 = 2 sec delay, Speed 10 = 0.1 sec delay.
class SimulationStatus(BaseModel):
status: str # "running", "paused", "stopped"
current_index: int
total_transactions: int
speed: float
class SimulationManager:
def __init__(self):
self.dataset: Optional[pd.DataFrame] = None
self.current_index = 0
self.is_running = False
self.speed = 1.0 # Default 1x speed
self.delay = 1.0 # Seconds between transactions
self.dataset_path = None
def load_dataset(self, path: str):
self.dataset_path = path
# Check if file exists, if not try to find it
if not os.path.exists(path):
# Try finding it in common locations
possible_paths = [
path,
f"dataset/{path}",
f"../dataset/{path}",
"dataset/test_dataset.csv.gz",
"dataset/test_dataset.csv"
]
for p in possible_paths:
if os.path.exists(p):
self.dataset_path = p
break
if self.dataset_path and os.path.exists(self.dataset_path):
print(f"Loading simulation dataset from {self.dataset_path}")
# Load only necessary columns to save memory if needed
self.dataset = pd.read_csv(self.dataset_path)
# Ensure it's sorted by step
if 'step' in self.dataset.columns:
self.dataset = self.dataset.sort_values('step')
print(f"Loaded {len(self.dataset)} transactions for simulation")
else:
print("Simulation dataset not found!")
def start(self):
if self.dataset is None:
raise ValueError("Dataset not loaded")
self.is_running = True
return self.get_status()
def stop(self):
self.is_running = False
return self.get_status()
def reset(self):
self.current_index = 0
self.is_running = False
return self.get_status()
def set_speed(self, speed: float):
self.speed = max(0.1, min(speed, 10.0)) # Clamp between 0.1 and 10
# Formula: Higher speed = Lower delay.
# Base delay = 2 seconds (at 1x)
# Delay = 2 / speed
self.delay = 2.0 / self.speed
return self.get_status()
def get_status(self):
return {
"status": "running" if self.is_running else "paused",
"current_index": self.current_index,
"total_transactions": len(self.dataset) if self.dataset is not None else 0,
"speed": self.speed
}
async def stream_generator(self) -> AsyncGenerator[str, None]:
"""
Yields SSE formatted data:
data: {json_content}\n\n
"""
while True:
if not self.is_running:
await asyncio.sleep(0.5)
# yield Comment to keep connection alive
yield ": keep-alive\n\n"
continue
if self.dataset is None or self.current_index >= len(self.dataset):
self.is_running = False
yield f"data: {json.dumps({'event': 'finished'})}\\n\\n"
break
# Get current transaction
row = self.dataset.iloc[self.current_index]
transaction = row.to_dict()
# Convert numpy types to native python types for JSON serialization
for key, value in transaction.items():
if hasattr(value, 'item'):
transaction[key] = value.item()
# Create payload
payload = {
"transaction": transaction,
"index": self.current_index,
"total": len(self.dataset)
}
yield f"data: {json.dumps(payload)}\\n\\n"
self.current_index += 1
# Wait based on speed
await asyncio.sleep(self.delay)
simulation_manager = SimulationManager()
|