Upload alpha_factory/personas/expression_compiler.py with huggingface_hub

alpha_factory/personas/expression_compiler.py

@@ -0,0 +1,129 @@
+"""
+Expression Compiler — Persona 2 (Hybrid: Jinja + Tinyfish LLM)
+Converts Blueprint JSON → valid BRAIN expression string.
+95% handled by templates. LLM only for novel structures.
+"""
+from jinja2 import Environment, BaseLoader
+from ..infra.llm_client import LLMClient
+from ..schemas import Blueprint, Expression
+
+
+# ─── Jinja Templates for proven archetypes ───────────────────────────────
+TEMPLATES = {
+    "value_quality_blend": """
+{%- set comps = [] -%}
+{%- for c in bp.components -%}
+  {%- set _ = comps.append(c.weight|string ~ " * group_zscore(rank(ts_mean(" ~ c.fields[0] ~ ", " ~ c.horizon_days ~ ")), " ~ bp.neutralization.value ~ ")") -%}
+{%- endfor -%}
+{{ comps | join(" + ") }}
+""",
+    "multi_horizon_mr": """
+{%- set main = bp.components[0] -%}
+zscore(ts_rank({{ main.fields[0] }}, {{ main.horizon_days }})) * {{ main.weight }} + zscore(ts_rank({{ bp.components[1].fields[0] if bp.components|length > 1 else main.fields[0] }}, {{ bp.components[1].horizon_days if bp.components|length > 1 else 20 }})) * {{ bp.components[1].weight if bp.components|length > 1 else 0.3 }}
+""",
+    "vol_scaled_shock": """
+{%- set c = bp.components[0] -%}
+zscore(ts_delta({{ c.fields[0] }}, {{ c.horizon_days }}) / (ts_std({{ c.fields[0] }}, {{ c.horizon_days * 4 }}) + 0.001))
+""",
+    "intraday_mr_decay": """
+{%- set c = bp.components[0] -%}
+ts_decay_linear(zscore(rank({{ c.fields[0] }})), {{ bp.decay }})
+""",
+    "pead_revisions": """
+{%- set c = bp.components[0] -%}
+group_zscore(ts_delta({{ c.fields[0] }}, {{ c.horizon_days }}), {{ bp.neutralization.value }})
+""",
+    "fundamental_yield_composite": """
+{%- set comps = [] -%}
+{%- for c in bp.components -%}
+  {%- set _ = comps.append(c.weight|string ~ " * zscore(rank(" ~ c.fields[0] ~ "))") -%}
+{%- endfor -%}
+ts_decay_linear({{ comps | join(" + ") }}, {{ bp.decay }})
+""",
+}
+
+# Jinja environment
+_env = Environment(loader=BaseLoader())
+
+
+COMPILER_SYSTEM_PROMPT = """You are a BRAIN expression compiler. Your ONLY job is to convert 
+a factor blueprint into a valid WorldQuant BRAIN expression.
+
+RULES:
+1. Use ONLY operators from the BRAIN operator catalogue.
+2. Every additive operand MUST be wrapped in zscore(), rank(), or group_zscore().
+3. Output ONLY the expression string — no explanation, no markdown.
+4. The expression must be syntactically valid (balanced parentheses, correct arity).
+5. Use ts_decay_linear for smoothing if decay > 0.
+6. Use group_zscore or indneutralize for neutralization within the expression if specified.
+
+Available operators: rank, zscore, group_zscore, group_rank, ts_mean, ts_std, ts_sum,
+ts_delta, ts_decay_linear, ts_rank, ts_argmax, ts_argmin, ts_correlation,
+ts_covariance, ts_regression, winsorize, abs, log, sign, power, sqrt, max, min,
+if_else, less, greater, filter, trade_when, ts_backfill, indneutralize
+"""
+
+
+async def compile_expression(
+    blueprint: Blueprint,
+    llm: LLMClient,
+    model: str | None = None,
+) -> Expression:
+    """
+    Convert a Blueprint to a BRAIN expression.
+    Uses Jinja template if archetype is known; LLM fallback for novel structures.
+    """
+    # Try template first (95% of cases)
+    if blueprint.archetype in TEMPLATES:
+        template_str = TEMPLATES[blueprint.archetype]
+        template = _env.from_string(template_str)
+        expr_text = template.render(bp=blueprint).strip()
+
+        # Extract metadata
+        fields_used = []
+        ops_used = []
+        for c in blueprint.components:
+            fields_used.extend(c.fields)
+            ops_used.extend(c.operators)
+
+        return Expression(
+            expression=expr_text,
+            fields_used=list(set(fields_used)),
+            operators_used=list(set(ops_used)),
+            archetype_used=blueprint.archetype,
+        )
+
+    # LLM fallback for novel archetypes
+    user_prompt = f"""Convert this factor blueprint into a valid BRAIN expression:
+
+Blueprint:
+- Theme: {blueprint.theme}
+- Components:
+{_format_components(blueprint)}
+- Neutralization: {blueprint.neutralization.value}
+- Decay: {blueprint.decay}
+- Target: cross-sectional rank → long/short
+
+Output a valid BRAIN expression. Wrap all additive operands in zscore() or rank().
+Apply ts_decay_linear with decay={blueprint.decay} if decay > 0."""
+
+    result = await llm.generate_json(
+        prompt=user_prompt,
+        schema=Expression,
+        model=model or llm.config.tinyfish_model,
+        temperature=0.1,
+        system_prompt=COMPILER_SYSTEM_PROMPT,
+    )
+
+    return result
+
+
+def _format_components(bp: Blueprint) -> str:
+    """Format components for the LLM prompt."""
+    lines = []
+    for i, c in enumerate(bp.components):
+        lines.append(
+            f"  {i+1}. {c.name}: fields={c.fields}, operators={c.operators}, "
+            f"horizon={c.horizon_days}d, weight={c.weight}, sign={c.sign_direction}"
+        )
+    return "\n".join(lines)