yuqei
128 lines
3.9 KiB
Python
128 lines
3.9 KiB
Python
from pathlib import Path
|
|
import json
|
|
import os
|
|
import sys
|
|
from typing import Iterable, Optional
|
|
|
|
from sqlalchemy.exc import SQLAlchemyError
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
if str(ROOT) not in sys.path:
|
|
sys.path.insert(0, str(ROOT))
|
|
|
|
from src.services.training import CaseGenerationResult, TrainingCaseGenerator
|
|
|
|
|
|
def run_training_generation(
|
|
generator: TrainingCaseGenerator,
|
|
instrument_id: int = 1,
|
|
timeframe: str = "1m",
|
|
case_types: Optional[Iterable[str]] = None,
|
|
limit_per_type: int = 10,
|
|
output_json: bool = False,
|
|
) -> int:
|
|
case_type_list = list(case_types) if case_types is not None else ["sweep", "fvg", "range", "mss_displacement"]
|
|
result = generator.generate_cases(
|
|
instrument_id=instrument_id,
|
|
timeframe=timeframe,
|
|
case_types=case_type_list,
|
|
limit_per_type=limit_per_type,
|
|
)
|
|
payload = _generation_payload(
|
|
result=result,
|
|
instrument_id=instrument_id,
|
|
timeframe=timeframe,
|
|
limit_per_type=limit_per_type,
|
|
case_types=case_type_list,
|
|
)
|
|
|
|
if output_json:
|
|
print(json.dumps(payload, default=str, ensure_ascii=False, sort_keys=True))
|
|
else:
|
|
print(f"Training cases generated: {result.generated}")
|
|
print(f"Training case ids: {', '.join(str(case_id) for case_id in result.case_ids) or 'none'}")
|
|
print(f"Training cases by type: {json.dumps(result.by_type, ensure_ascii=False, sort_keys=True)}")
|
|
return 0
|
|
|
|
|
|
def main(env=None) -> int:
|
|
env = os.environ if env is None else env
|
|
output_json = _get_bool(env, "TRAINING_CASES_OUTPUT_JSON", False)
|
|
instrument_id = _get_int(env, "TRAINING_INSTRUMENT_ID", 1)
|
|
timeframe = str(env.get("TRAINING_TIMEFRAME") or "1m")
|
|
limit_per_type = _get_int(env, "TRAINING_LIMIT_PER_TYPE", 10)
|
|
case_types = _get_case_types(env, "TRAINING_CASE_TYPES")
|
|
|
|
try:
|
|
return run_training_generation(
|
|
generator=TrainingCaseGenerator(),
|
|
instrument_id=instrument_id,
|
|
timeframe=timeframe,
|
|
case_types=case_types,
|
|
limit_per_type=limit_per_type,
|
|
output_json=output_json,
|
|
)
|
|
except SQLAlchemyError as exc:
|
|
payload = {
|
|
"passed": False,
|
|
"instrument_id": instrument_id,
|
|
"timeframe": timeframe,
|
|
"limit_per_type": limit_per_type,
|
|
"case_types": case_types or ["sweep", "fvg", "range", "mss_displacement"],
|
|
"generated": 0,
|
|
"by_type": {},
|
|
"case_ids": [],
|
|
"error_type": exc.__class__.__name__,
|
|
"error": str(exc),
|
|
}
|
|
if output_json:
|
|
print(json.dumps(payload, default=str, ensure_ascii=False, sort_keys=True))
|
|
else:
|
|
print(f"Database unavailable for training case generation: {exc.__class__.__name__}")
|
|
return 2
|
|
|
|
|
|
def _generation_payload(
|
|
result: CaseGenerationResult,
|
|
instrument_id: int,
|
|
timeframe: str,
|
|
limit_per_type: int,
|
|
case_types: list[str],
|
|
) -> dict:
|
|
return {
|
|
"passed": True,
|
|
"instrument_id": instrument_id,
|
|
"timeframe": timeframe,
|
|
"limit_per_type": limit_per_type,
|
|
"case_types": case_types,
|
|
"generated": result.generated,
|
|
"by_type": result.by_type,
|
|
"case_ids": result.case_ids,
|
|
}
|
|
|
|
|
|
def _get_int(env, key: str, default: int) -> int:
|
|
value = env.get(key)
|
|
if value is None or value == "":
|
|
return default
|
|
return int(value)
|
|
|
|
|
|
def _get_bool(env, key: str, default: bool) -> bool:
|
|
value = env.get(key)
|
|
if value is None or value == "":
|
|
return default
|
|
return str(value).strip().lower() in {"1", "true", "yes", "on"}
|
|
|
|
|
|
def _get_case_types(env, key: str) -> Optional[list[str]]:
|
|
value = env.get(key)
|
|
if value is None or str(value).strip() == "":
|
|
return None
|
|
case_types = [part.strip() for part in str(value).split(",") if part.strip()]
|
|
return case_types or None
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|