cnhkmcp 2.3.0__py3-none-any.whl → 2.3.1__py3-none-any.whl

cnhkmcp/untracked/APP/trailSomeAlphas/run_pipeline.py

@@ -3,6 +3,7 @@ import datetime as dt
 import json
 import os
 import re
+import shutil
 import subprocess
 import sys
 import csv
@@ -327,6 +328,7 @@ def build_prompt(
     region: str,
     delay: int,
     universe: str,
+    data_type: str,
     fields_summary: list[dict],
     field_count: int,
     feature_engineering_skill_md: str,
@@ -336,8 +338,7 @@ def build_prompt(
 ):
     # NOTE: The user requested that we DO NOT invent our own system prompt.
     # Instead, we embed the two skill specs as the authoritative instructions.
-    system_prompt = "\n".join(
-        [
+    prompt_lines = [
             "You are executing two skills in sequence:",
             "1) brain-data-feature-engineering",
             "2) brain-feature-implementation",
@@ -353,6 +354,15 @@ def build_prompt(
             "-------",
             f'"allowed_placeholders": {allowed_metric_suffixes}',
             "",
+        ]
+
+    if str(data_type).upper() == "VECTOR":
+        prompt_lines.append(
+            "since all the following the data is vector type data, before you do any process, you should choose a vector operator to generate its statistical feature to use, the data cannot be directly use. for example, if datafieldA and datafieldB are vector type data, you can use vec_avg(datafieldA) -  vec_avg(datafieldB), where vec_avg() operator is used to generate the average of the data on a certain date. similarly, vector type operator can only be used on the vector type operator directly and cannot be nested, for example vec_avg(vec_sum(datafield)) is a false use."
+        )
+
+    prompt_lines.extend(
+        [
             "CRITICAL OUTPUT RULES (to ensure implement_idea.py can generate expressions):",
             "- Every Implementation Example MUST be a Python format template using {variable}.",
             "- Every {variable} MUST come from the allowed_placeholders list provided in user content.",
@@ -366,6 +376,8 @@ def build_prompt(
         ]
     )
 
+    system_prompt = "\n".join(prompt_lines)
+
     user_prompt = {
         "instructions": {
             "output_format": "Fill OUTPUT_TEMPLATE.md with concrete content.",
@@ -732,6 +744,21 @@ def run_script(args_list: list[str], cwd: Path):
         )
     return result.stdout
 
+
+def delete_path_if_exists(path: Path):
+    """Best-effort delete a file or directory."""
+
+    try:
+        if not path.exists():
+            return
+        if path.is_dir():
+            shutil.rmtree(path, ignore_errors=True)
+        else:
+            path.unlink(missing_ok=True)
+    except Exception:
+        # Best-effort cleanup only; rerun should still proceed.
+        return
+
 def main():
     parser = argparse.ArgumentParser(description="Run feature engineering + implementation pipeline")
     parser.add_argument("--data-category", required=True, help="Dataset category (e.g., analyst, fundamental)")
@@ -740,6 +767,12 @@ def main():
     parser.add_argument("--universe", default="TOP3000", help="Universe (default: TOP3000)")
     parser.add_argument("--dataset-id", required=True, help="Dataset id (required)")
     parser.add_argument("--instrument-type", default="EQUITY", help="Instrument type (default: EQUITY)")
+    parser.add_argument(
+        "--data-type",
+        default="MATRIX",
+        choices=["MATRIX", "VECTOR"],
+        help="Data type to request from BRAIN datafields (MATRIX or VECTOR). Default: MATRIX",
+    )
     parser.add_argument("--ideas-file", default=None, help="Use existing ideas markdown instead of generating")
     parser.add_argument(
         "--regen-ideas",
@@ -774,96 +807,105 @@ def main():
     email, password = load_brain_credentials_from_env_or_args(args.username, args.password, config_path)
     session = start_brain_session(email, password)
 
+    # Always rerun cleanly: remove prior generated artifacts so we never reuse stale ideas/data.
+    # - If --ideas-file is provided, we treat it as user-managed input and do NOT delete it.
+    # - We DO delete the dataset-specific folder under feature-implementation/data.
+    if not args.ideas_file:
+        default_ideas = (
+            FEATURE_ENGINEERING_DIR
+            / "output_report"
+            / f"{args.region}_delay{args.delay}_{args.dataset_id}_ideas.md"
+        )
+        delete_path_if_exists(default_ideas)
+
+    guessed_dataset_folder = f"{safe_dataset_id(args.dataset_id)}_{args.region}_delay{args.delay}"
+    guessed_dataset_dir = FEATURE_IMPLEMENTATION_DIR / "data" / guessed_dataset_folder
+    delete_path_if_exists(guessed_dataset_dir)
+
     ideas_path = None
     if args.ideas_file:
         ideas_path = Path(args.ideas_file).resolve()
         if not ideas_path.exists():
             raise FileNotFoundError(f"Ideas file not found: {ideas_path}")
     else:
-        default_ideas = (
-            FEATURE_ENGINEERING_DIR
-            / "output_report"
-            / f"{args.region}_delay{args.delay}_{args.dataset_id}_ideas.md"
+        # Always regenerate ideas (never reuse an existing markdown report).
+        datasets_df = ace_lib.get_datasets(
+            session,
+            instrument_type=args.instrument_type,
+            region=args.region,
+            delay=args.delay,
+            universe=args.universe,
+            theme="ALL",
         )
-        if default_ideas.exists() and not args.regen_ideas:
-            ideas_path = default_ideas
-        else:
-            datasets_df = ace_lib.get_datasets(
-                session,
-                instrument_type=args.instrument_type,
-                region=args.region,
-                delay=args.delay,
-                universe=args.universe,
-                theme="ALL",
-            )
-
-            dataset_name = None
-            dataset_description = None
-            id_col = pick_first_present_column(datasets_df, ["id", "dataset_id", "datasetId"])
-            name_col = pick_first_present_column(datasets_df, ["name", "dataset_name", "datasetName"])
-            desc_col = pick_first_present_column(datasets_df, ["description", "desc", "dataset_description"])
-            if id_col:
-                matched = datasets_df[datasets_df[id_col].astype(str) == str(args.dataset_id)]
-                if not matched.empty:
-                    row = matched.iloc[0]
-                    dataset_name = row.get(name_col) if name_col else None
-                    dataset_description = row.get(desc_col) if desc_col else None
-
-            fields_df = ace_lib.get_datafields(
-                session,
-                instrument_type=args.instrument_type,
-                region=args.region,
-                delay=args.delay,
-                universe=args.universe,
-                dataset_id=args.dataset_id,
-                data_type="ALL",
-            )
-
-            fields_summary, field_count = build_field_summary(fields_df, max_fields=args.max_fields)
-
-            feature_engineering_skill_md = read_text_optional(FEATURE_ENGINEERING_DIR / "SKILL.md")
-            feature_implementation_skill_md = read_text_optional(FEATURE_IMPLEMENTATION_DIR / "SKILL.md")
-            allowed_metric_suffixes = build_allowed_metric_suffixes(fields_df, max_suffixes=300)
-
-            allowed_operators = []
-            if not args.no_operators_in_prompt:
-                try:
-                    operators_df = ace_lib.get_operators(session)
-                    keep_vector = _vector_ratio_from_datafields_df(fields_df) > 0.5
-                    _, allowed_ops, _ = filter_operators_df(operators_df, keep_vector=keep_vector)
-                    if args.max_operators is not None and args.max_operators > 0:
-                        allowed_operators = allowed_ops[: args.max_operators]
-                    else:
-                        allowed_operators = allowed_ops
-                except Exception as exc:
-                    print(f"Warning: failed to fetch/filter operators; continuing without operators in prompt. Error: {exc}", file=sys.stderr)
-
-            system_prompt, user_prompt = build_prompt(
-                dataset_id=args.dataset_id,
-                dataset_name=dataset_name,
-                dataset_description=dataset_description,
-                data_category=args.data_category,
-                region=args.region,
-                delay=args.delay,
-                universe=args.universe,
-                fields_summary=fields_summary,
-                field_count=field_count,
-                feature_engineering_skill_md=feature_engineering_skill_md,
-                feature_implementation_skill_md=feature_implementation_skill_md,
-                allowed_metric_suffixes=allowed_metric_suffixes,
-                allowed_operators=allowed_operators,
-            )
-            
-            api_key = (
-                args.moonshot_api_key
-                or os.environ.get("MOONSHOT_API_KEY")
-            )
-            if not api_key:
-                raise ValueError("Moonshot API key missing. Set MOONSHOT_API_KEY or pass --moonshot-api-key")
-
-            report = call_moonshot(api_key, args.moonshot_model, system_prompt, user_prompt)
-            # Save first, then normalize placeholders after dataset download.
-            ideas_path = save_ideas_report(report, args.region, args.delay, args.dataset_id)
+
+        dataset_name = None
+        dataset_description = None
+        id_col = pick_first_present_column(datasets_df, ["id", "dataset_id", "datasetId"])
+        name_col = pick_first_present_column(datasets_df, ["name", "dataset_name", "datasetName"])
+        desc_col = pick_first_present_column(datasets_df, ["description", "desc", "dataset_description"])
+        if id_col:
+            matched = datasets_df[datasets_df[id_col].astype(str) == str(args.dataset_id)]
+            if not matched.empty:
+                row = matched.iloc[0]
+                dataset_name = row.get(name_col) if name_col else None
+                dataset_description = row.get(desc_col) if desc_col else None
+
+        fields_df = ace_lib.get_datafields(
+            session,
+            instrument_type=args.instrument_type,
+            region=args.region,
+            delay=args.delay,
+            universe=args.universe,
+            dataset_id=args.dataset_id,
+            data_type=args.data_type,
+        )
+
+        fields_summary, field_count = build_field_summary(fields_df, max_fields=args.max_fields)
+
+        feature_engineering_skill_md = read_text_optional(FEATURE_ENGINEERING_DIR / "SKILL.md")
+        feature_implementation_skill_md = read_text_optional(FEATURE_IMPLEMENTATION_DIR / "SKILL.md")
+        allowed_metric_suffixes = build_allowed_metric_suffixes(fields_df, max_suffixes=300)
+
+        allowed_operators = []
+        if not args.no_operators_in_prompt:
+            try:
+                operators_df = ace_lib.get_operators(session)
+                keep_vector = _vector_ratio_from_datafields_df(fields_df) > 0.5
+                _, allowed_ops, _ = filter_operators_df(operators_df, keep_vector=keep_vector)
+                if args.max_operators is not None and args.max_operators > 0:
+                    allowed_operators = allowed_ops[: args.max_operators]
+                else:
+                    allowed_operators = allowed_ops
+            except Exception as exc:
+                print(f"Warning: failed to fetch/filter operators; continuing without operators in prompt. Error: {exc}", file=sys.stderr)
+
+        system_prompt, user_prompt = build_prompt(
+            dataset_id=args.dataset_id,
+            dataset_name=dataset_name,
+            dataset_description=dataset_description,
+            data_category=args.data_category,
+            region=args.region,
+            delay=args.delay,
+            universe=args.universe,
+            data_type=args.data_type,
+            fields_summary=fields_summary,
+            field_count=field_count,
+            feature_engineering_skill_md=feature_engineering_skill_md,
+            feature_implementation_skill_md=feature_implementation_skill_md,
+            allowed_metric_suffixes=allowed_metric_suffixes,
+            allowed_operators=allowed_operators,
+        )
+
+        api_key = (
+            args.moonshot_api_key
+            or os.environ.get("MOONSHOT_API_KEY")
+        )
+        if not api_key:
+            raise ValueError("Moonshot API key missing. Set MOONSHOT_API_KEY or pass --moonshot-api-key")
+
+        report = call_moonshot(api_key, args.moonshot_model, system_prompt, user_prompt)
+        # Save first, then normalize placeholders after dataset download.
+        ideas_path = save_ideas_report(report, args.region, args.delay, args.dataset_id)
 
     ideas_text = ideas_path.read_text(encoding="utf-8")
 
@@ -891,11 +933,19 @@ def main():
             args.universe,
             "--instrument-type",
             args.instrument_type,
+            "--data-type",
+            args.data_type,
         ],
         cwd=FEATURE_IMPLEMENTATION_SCRIPTS,
     )
 
     dataset_folder = f"{safe_dataset_id(dataset_id)}_{args.region}_delay{args.delay}"
+
+    # If the ideas file references a different dataset id than the CLI args,
+    # ensure we also clean that dataset folder before fetching.
+    if dataset_folder != guessed_dataset_folder:
+        delete_path_if_exists(FEATURE_IMPLEMENTATION_DIR / "data" / dataset_folder)
+
     dataset_csv_path = FEATURE_IMPLEMENTATION_DIR / "data" / dataset_folder / f"{dataset_folder}.csv"
     if not dataset_csv_path.exists():
         raise RuntimeError(