synth-ai 0.2.4.dev5__py3-none-any.whl → 0.2.4.dev6__py3-none-any.whl

synth_ai/learning/prompts/run_random_search_banking77.py

@@ -0,0 +1,305 @@
+"""
+Example: Random Search optimizer on Banking77 using Groq gpt-oss-20b.
+
+Requires:
+- .env with GROQ_API_KEY
+- datasets (`uv add datasets` if needed)
+
+Run:
+- uv run -q python -m synth_ai.learning.prompts.run_random_search_banking77
+"""
+
+from __future__ import annotations
+
+import asyncio
+import os
+import random
+from dataclasses import dataclass, replace
+from types import SimpleNamespace
+from tqdm import tqdm
+from typing import Any, Dict, List, Sequence, Tuple
+
+from dotenv import load_dotenv
+from datasets import load_dataset
+
+from synth_ai.lm.core.main_v3 import LM, build_messages
+import json
+import time
+from pathlib import Path
+from synth_ai.learning.prompts.random_search import random_search_compile
+
+
+def choose_label(pred: str, label_names: List[str]) -> str:
+    norm = (pred or "").strip().lower()
+    d = {ln.lower(): ln for ln in label_names}
+    if norm in d:
+        return d[norm]
+    def score(cand: str) -> int:
+        c = cand.lower()
+        return sum(1 for w in c.split() if w in norm)
+    return max(label_names, key=score)
+
+
+def accuracy(pred: str, gold: str, labels: List[str]) -> float:
+    return 1.0 if choose_label(pred, labels) == gold else 0.0
+
+
+@dataclass
+class StudentProgram:
+    lm: LM
+    label_names: List[str]
+    instruction: str
+    demos: List[Tuple[str, str]]
+
+    def reset_copy(self):
+        return replace(self, instruction=self.instruction, demos=list(self.demos))
+
+    def deepcopy(self):
+        return replace(self, instruction=str(self.instruction), demos=list(self.demos))
+
+    def with_demos(self, demos: List[Tuple[str, str]]):
+        return replace(self, demos=list(demos))
+
+    def run(self, x: str) -> str:
+        # Build a prompt with optional demos
+        examples = "\n".join(f"Input: {a}\nLabel: {b}" for a, b in self.demos)
+        sys = self.instruction or "You are an intent classifier for Banking77."
+        user = (f"Examples:\n{examples}\n\n" if examples else "") + f"Message: {x}\nLabel:"
+        messages = build_messages(sys, user, images_bytes=None, model_name=self.lm.model)
+        # Call LM synchronously via asyncio
+        async def _call():
+            resp = await self.lm.respond_async(messages=messages)
+            return (resp.raw_response or "").strip()
+        return asyncio.run(_call())
+
+    async def _apredict(self, x: str):
+        examples = "\n".join(f"Input: {a}\nLabel: {b}" for a, b in self.demos)
+        sys = self.instruction or "You are an intent classifier for Banking77."
+        user = (f"Examples:\n{examples}\n\n" if examples else "") + f"Message: {x}\nLabel:"
+        messages = build_messages(sys, user, images_bytes=None, model_name=self.lm.model)
+        resp = await self.lm.respond_async(messages=messages)
+        return (resp.raw_response or "").strip(), (resp.usage or {})
+
+
+def main():
+    load_dotenv()
+    random.seed(0)
+
+    model = os.getenv("MODEL", "openai/gpt-oss-20b")
+    vendor = os.getenv("VENDOR", "groq")
+    lm = LM(model=model, vendor=vendor, temperature=0.0)
+
+    print("Loading Banking77 dataset (train/dev split of test for demo)...")
+    ds = load_dataset("banking77")
+    label_names: List[str] = ds["test"].features["label"].names  # type: ignore
+
+    # Create small train/val from the test split for speed
+    all_items = [(r["text"], label_names[int(r["label"])]) for r in ds["test"]]
+    random.shuffle(all_items)
+    trainset: Sequence[Tuple[str, str]] = all_items[:40]
+    valset: Sequence[Tuple[str, str]] = all_items[40:60]  # 20 examples
+
+    student = StudentProgram(
+        lm=lm,
+        label_names=label_names,
+        instruction="You are an intent classifier for the Banking77 dataset. Return exactly one label.",
+        demos=[],
+    )
+
+    def metric(yhat: str, y: str) -> float:
+        return accuracy(yhat, y, label_names)
+
+    total_candidates = 3 + 3  # zero-shot, labeled few-shot, bootstrapped + 3 random seeds
+    print(f"Running Random Search optimizer ({total_candidates} candidates, parallel eval of 20 questions)...")
+
+    def eval_parallel(program: StudentProgram, dataset: Sequence[Tuple[str, str]], metric_fn):
+        async def _run():
+            xs = [x for x, _ in dataset]
+            ys = [y for _, y in dataset]
+            preds: List[Optional[str]] = [None] * len(xs)
+            sem = asyncio.Semaphore(int(os.getenv("CONCURRENCY", "5")))
+
+            async def worker(i: int, x: str, y: str):
+                import time
+                t_start = time.monotonic()
+                try:
+                    async with sem:
+                        pred, usage = await asyncio.wait_for(
+                            program._apredict(x),
+                            timeout=float(os.getenv("TIMEOUT_S", "45")),
+                        )
+                        t_end = time.monotonic()
+                        return i, y, pred, t_start, t_end, usage or {}
+                except asyncio.CancelledError:
+                    # Respect cancellation but return a placeholder record so scheduler can proceed
+                    t_end = time.monotonic()
+                    return i, y, "", t_start, t_end, {}
+                except Exception:
+                    t_end = time.monotonic()
+                    return i, y, "", t_start, t_end, {}
+
+            tasks = [asyncio.create_task(worker(i, x, y)) for i, (x, y) in enumerate(zip(xs, ys))]
+            correct_sum = 0.0
+            processed = 0
+            import time, statistics
+            durations: List[float] = []
+            in_tok_sum = 0
+            out_tok_sum = 0
+            in_tok_count = 0
+            out_tok_count = 0
+            details: List[Dict[str, Any]] = []
+            t_batch_start = time.monotonic()
+            deadline = float(os.getenv("BATCH_DEADLINE_S", "20"))
+            with tqdm(total=len(tasks), desc="Rollouts", leave=False) as pbar:
+                pending = set(tasks)
+                # Process completions until all done or deadline reached
+                while pending:
+                    elapsed = time.monotonic() - t_batch_start
+                    remaining = max(0.0, deadline - elapsed)
+                    if remaining <= 0.0:
+                        # Cancel any remaining
+                        for t in pending:
+                            t.cancel()
+                        done, _ = await asyncio.wait(pending, return_when=asyncio.ALL_COMPLETED)
+                        # Record canceled as zeros
+                        for task in done:
+                            try:
+                                i, y_true, pred, t_start, t_end, usage = task.result()
+                            except Exception:
+                                # Unknown index: we can't recover; skip as it's canceled before start
+                                continue
+                            # Already processed ones shouldn't be in pending; skip
+                        break
+                    # Wait for at least one completion within remaining time (polling granularity <= 1s)
+                    timeout = min(1.0, remaining)
+                    done, pending = await asyncio.wait(pending, timeout=timeout, return_when=asyncio.FIRST_COMPLETED)
+                    for task in done:
+                        try:
+                            i, y_true, pred, t_start, t_end, usage = task.result()
+                        except BaseException:
+                            # Treat as failure/cancelled
+                            continue
+                        durations.append(max(0.0, t_end - t_start))
+                        preds[i] = pred
+                        processed += 1
+                        try:
+                            correct_sum += float(metric_fn(pred, y_true))
+                        except Exception:
+                            pass
+                        try:
+                            pt = usage.get("prompt_tokens") or usage.get("input_tokens")
+                            ct = usage.get("completion_tokens") or usage.get("output_tokens")
+                            if isinstance(pt, (int, float)):
+                                in_tok_sum += int(pt)
+                                in_tok_count += 1
+                            if isinstance(ct, (int, float)):
+                                out_tok_sum += int(ct)
+                                out_tok_count += 1
+                        except Exception:
+                            pass
+                        details.append({
+                            "index": i,
+                            "seconds": max(0.0, t_end - t_start),
+                            "score": float(metric_fn(pred, y_true)),
+                            "usage": {
+                                "prompt_tokens": usage.get("prompt_tokens") or usage.get("input_tokens"),
+                                "completion_tokens": usage.get("completion_tokens") or usage.get("output_tokens"),
+                            },
+                        })
+                        pbar.update(1)
+                        med = statistics.median(durations) if durations else 0.0
+                        mx = max(durations) if durations else 0.0
+                        avg_in = (in_tok_sum / in_tok_count) if in_tok_count else 0.0
+                        avg_out = (out_tok_sum / out_tok_count) if out_tok_count else 0.0
+                        pbar.set_postfix({
+                            "acc": f"{(correct_sum/processed):.2f}",
+                            "done": f"{processed}/{len(tasks)}",
+                            "med_s": f"{med:.1f}",
+                            "max_s": f"{mx:.1f}",
+                            "tin": f"{avg_in:.1f}",
+                            "tout": f"{avg_out:.1f}",
+                        })
+            # Compute score only from completed/successful rollouts (drop timeouts/cancelled)
+            subs = [float(d.get("score", 0.0)) for d in details]
+            result = SimpleNamespace(score=(sum(subs) / max(1, len(subs))), subscores=subs)
+            result.details = details
+            result.mean_in = (in_tok_sum / in_tok_count) if in_tok_count else 0.0
+            result.mean_out = (out_tok_sum / out_tok_count) if out_tok_count else 0.0
+            return result
+        return asyncio.run(_run())
+    pbar = tqdm(total=total_candidates, desc="Candidates")
+    candidate_eval_details: Dict[int, Any] = {}
+    def on_cand(idx: int, score: float, res, intervention):
+        pbar.update(1)
+        pbar.set_postfix({"score": f"{score:.2f}"})
+        # store per-instance details (for apples-to-apples)
+        try:
+            candidate_eval_details[idx] = {
+                "score": score,
+                "mean_in": getattr(res, "mean_in", None),
+                "mean_out": getattr(res, "mean_out", None),
+                "instances": getattr(res, "details", None),
+            }
+        except Exception:
+            pass
+        # visible summary line per candidate
+        kind = intervention.get("kind", "candidate") if isinstance(intervention, dict) else "candidate"
+        label = intervention.get("label") if isinstance(intervention, dict) else None
+        seed = intervention.get("seed") if isinstance(intervention, dict) else None
+        processed = len(getattr(res, "details", []) or [])
+        from tqdm import tqdm as _tqdm
+        _tqdm.write(
+            f"Candidate {idx}/{total_candidates} [{kind}{'' if label is None else f', label={label}'}{'' if seed is None else f', seed={seed}'}]: "
+            f"score={score:.2f} | mean tin/tout={getattr(res, 'mean_in', 0):.1f}/{getattr(res, 'mean_out', 0):.1f} | N={processed}"
+        )
+
+    best, records = random_search_compile(
+        student=student,
+        trainset=trainset,
+        valset=valset,
+        metric=metric,
+        evaluate_fn=eval_parallel,
+        max_bootstrapped_demos=0,
+        max_labeled_demos=4,
+        max_rounds=2,
+        num_candidate_programs=3,
+        on_candidate_evaluated=on_cand,
+    )
+    pbar.close()
+
+    # Evaluate best on holdout (valset) with parallel rollouts
+    print("Evaluating best program on val (parallel rollouts)...")
+    best_res = eval_parallel(best, valset, metric)
+    correct = int(round(best_res.score * max(1, len(best_res.subscores))))
+    print(
+        "Best program accuracy on val: "
+        f"{correct}/{len(valset)} ({best_res.score:.2%}) "
+        f"| mean tokens in/out: {getattr(best_res, 'mean_in', 0):.1f}/{getattr(best_res, 'mean_out', 0):.1f}"
+    )
+
+    # Save per-candidate scores and interventions
+    out = {
+        "context": {
+            "model": model,
+            "vendor": vendor,
+            "train_size": len(trainset),
+            "val_size": len(valset),
+        },
+        "candidates": records,
+        "candidate_eval_details": candidate_eval_details,
+        "best_eval_details": {
+            "score": best_res.score,
+            "mean_in": getattr(best_res, "mean_in", None),
+            "mean_out": getattr(best_res, "mean_out", None),
+            "instances": getattr(best_res, "details", None),
+        },
+    }
+    out_dir = Path(__file__).parent
+    fname = str(out_dir / f"random_search_banking77_{int(time.time())}.json")
+    with open(fname, "w") as f:
+        json.dump(out, f, indent=2)
+    print(f"Saved candidate records to {fname}")
+
+
+if __name__ == "__main__":
+    main()

synth_ai/lm/injection.py

@@ -0,0 +1,81 @@
+from __future__ import annotations
+
+import contextvars
+from contextlib import contextmanager
+from typing import Any, Dict, List, Optional
+
+Rule = Dict[str, Any]
+
+_rules_ctx: contextvars.ContextVar[Optional[List[Rule]]] = contextvars.ContextVar(
+    "injection_rules", default=None
+)
+
+
+def set_injection_rules(rules: List[Rule]):
+    """Set prompt-injection rules for the current context and return a reset token.
+
+    Each rule must be a dict with at least keys: "find" and "replace" (strings).
+    Optional: "roles" as a list of role names to scope the replacement.
+    """
+    if not isinstance(rules, list) or not all(
+        isinstance(r, dict) and "find" in r and "replace" in r for r in rules
+    ):
+        raise ValueError("Injection rules must be a list of dicts with 'find' and 'replace'")
+    return _rules_ctx.set(rules)
+
+
+def get_injection_rules() -> Optional[List[Rule]]:
+    """Get the current context's injection rules, if any."""
+    return _rules_ctx.get()
+
+
+def clear_injection_rules(token) -> None:
+    """Reset the injection rules to the previous value using the provided token."""
+    _rules_ctx.reset(token)
+
+
+def apply_injection(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+    """Apply ordered substring replacements to text parts of messages in place.
+
+    - Only modifies `str` content or list parts where `part["type"] == "text"`.
+    - Honors optional `roles` scoping in each rule.
+    - Returns the input list for convenience.
+    """
+    rules = get_injection_rules()
+    if not rules:
+        return messages
+
+    for m in messages:
+        role = m.get("role")
+        content = m.get("content")
+        if isinstance(content, str):
+            new_content = content
+            for r in rules:
+                allowed_roles = r.get("roles")
+                if allowed_roles is not None and role not in allowed_roles:
+                    continue
+                new_content = new_content.replace(str(r["find"]), str(r["replace"]))
+            m["content"] = new_content
+        elif isinstance(content, list):
+            for part in content:
+                if part.get("type") == "text":
+                    text = part.get("text", "")
+                    new_text = text
+                    for r in rules:
+                        allowed_roles = r.get("roles")
+                        if allowed_roles is not None and role not in allowed_roles:
+                            continue
+                        new_text = new_text.replace(str(r["find"]), str(r["replace"]))
+                    part["text"] = new_text
+    return messages
+
+
+@contextmanager
+def injection_rules_ctx(rules: List[Rule]):
+    """Context manager to temporarily apply injection rules within the block."""
+    tok = set_injection_rules(rules)
+    try:
+        yield
+    finally:
+        clear_injection_rules(tok)
+

synth_ai/lm/overrides.py

@@ -0,0 +1,204 @@
+from __future__ import annotations
+
+from contextlib import contextmanager
+from typing import Any, Dict, List, Optional, Tuple
+import contextvars
+
+from synth_ai.lm.injection import (
+    set_injection_rules,
+    clear_injection_rules,
+    apply_injection as _apply_injection,
+)
+
+# Context to hold a list of override specs to evaluate per-call
+# Each spec shape (minimal v1):
+# {
+#   "match": {"contains": "atm", "role": "user" | "system" | None},
+#   "injection_rules": [{"find": str, "replace": str, "roles": Optional[List[str]]}],
+#   "params": { ... api params to override ... },
+#   "tools": { ... optional tools overrides ... },
+# }
+_override_specs_ctx: contextvars.ContextVar[Optional[List[Dict[str, Any]]]] = contextvars.ContextVar(
+    "override_specs", default=None
+)
+
+# ContextVars actually applied for the specific call once matched
+_param_overrides_ctx: contextvars.ContextVar[Optional[Dict[str, Any]]] = contextvars.ContextVar(
+    "param_overrides", default=None
+)
+_tool_overrides_ctx: contextvars.ContextVar[Optional[Dict[str, Any]]] = contextvars.ContextVar(
+    "tool_overrides", default=None
+)
+_current_override_label_ctx: contextvars.ContextVar[Optional[str]] = contextvars.ContextVar(
+    "override_label", default=None
+)
+
+
+def set_override_specs(specs: List[Dict[str, Any]]):
+    if not isinstance(specs, list):
+        raise ValueError("override specs must be a list of dicts")
+    return _override_specs_ctx.set(specs)
+
+
+def get_override_specs() -> Optional[List[Dict[str, Any]]]:
+    return _override_specs_ctx.get()
+
+
+def clear_override_specs(token) -> None:
+    _override_specs_ctx.reset(token)
+
+
+def _matches(spec: Dict[str, Any], messages: List[Dict[str, Any]]) -> bool:
+    match = spec.get("match") or {}
+    contains = match.get("contains")
+    role = match.get("role")  # optional
+    if not contains:
+        # no match criteria means always apply
+        return True
+    contains_l = str(contains).lower()
+    for m in messages:
+        if role and m.get("role") != role:
+            continue
+        c = m.get("content")
+        if isinstance(c, str) and contains_l in c.lower():
+            return True
+        if isinstance(c, list):
+            for part in c:
+                if part.get("type") == "text" and contains_l in str(part.get("text", "")).lower():
+                    return True
+    return False
+
+
+def resolve_override_for_messages(messages: List[Dict[str, Any]]) -> Optional[Dict[str, Any]]:
+    specs = get_override_specs() or []
+    for spec in specs:
+        try:
+            if _matches(spec, messages):
+                return spec
+        except Exception:
+            # On matcher errors, skip spec
+            continue
+    return None
+
+
+def apply_injection(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+    # Delegate to injection.apply_injection
+    return _apply_injection(messages)
+
+
+def apply_param_overrides(api_params: Dict[str, Any]) -> Dict[str, Any]:
+    ov = _param_overrides_ctx.get()
+    if not ov:
+        return api_params
+    # Shallow merge only known keys users provided
+    for k, v in ov.items():
+        api_params[k] = v
+    return api_params
+
+
+def apply_tool_overrides(api_params: Dict[str, Any]) -> Dict[str, Any]:
+    """Apply tool overrides to OpenAI/Anthropic-like api_params in place.
+
+    Supports keys under spec["tools"]:
+      - set_tools: replace tools entirely
+      - add_tools: append tools
+      - remove_tools_by_name: remove by function name
+      - tool_choice: set tool_choice param
+    """
+    ov = _tool_overrides_ctx.get()
+    if not ov:
+        return api_params
+    tov = ov.get("tools") if isinstance(ov, dict) else None
+    if tov:
+        tools = api_params.get("tools")
+        if "set_tools" in tov:
+            tools = tov["set_tools"]
+        if "add_tools" in tov:
+            tools = (tools or []) + tov["add_tools"]
+        if "remove_tools_by_name" in tov and tools:
+            names = set(tov["remove_tools_by_name"])  # function names
+            new_tools = []
+            for t in tools:
+                try:
+                    # OpenAI dict style
+                    fn = t.get("function", {}).get("name") if isinstance(t, dict) else None
+                except Exception:
+                    fn = None
+                # If BaseTool objects slipped through
+                if fn is None:
+                    fn = getattr(t, "function_name", None)
+                if fn is None or fn not in names:
+                    new_tools.append(t)
+            tools = new_tools
+        if tools is not None:
+            api_params["tools"] = tools
+        if "tool_choice" in tov:
+            api_params["tool_choice"] = tov["tool_choice"]
+    return api_params
+
+
+@contextmanager
+def use_overrides_for_messages(messages: List[Dict[str, Any]]):
+    """Resolve an override spec against messages and apply its contexts within the scope.
+
+    - Sets injection rules and param overrides if present on the matched spec.
+    - Yields, then resets ContextVars to previous values.
+    """
+    spec = resolve_override_for_messages(messages) or {}
+    inj_rules = spec.get("injection_rules")
+    params = spec.get("params")
+    inj_tok = None
+    param_tok = None
+    tool_tok = None
+    label_tok = None
+    try:
+        if inj_rules:
+            inj_tok = set_injection_rules(inj_rules)
+        if params:
+            param_tok = _param_overrides_ctx.set(params)
+        tools = spec.get("tools")
+        if tools:
+            tool_tok = _tool_overrides_ctx.set({"tools": tools})
+        lbl = spec.get("label")
+        if lbl:
+            label_tok = _current_override_label_ctx.set(str(lbl))
+        yield
+    finally:
+        if inj_tok is not None:
+            clear_injection_rules(inj_tok)
+        if param_tok is not None:
+            _param_overrides_ctx.reset(param_tok)
+        if tool_tok is not None:
+            _tool_overrides_ctx.reset(tool_tok)
+        if label_tok is not None:
+            _current_override_label_ctx.reset(label_tok)
+
+
+def get_current_override_label() -> Optional[str]:
+    return _current_override_label_ctx.get()
+
+
+class LMOverridesContext:
+    """Context manager to register per-call override specs.
+
+    Usage:
+        with LMOverridesContext([
+            {"match": {"contains": "atm", "role": "user"}, "injection_rules": [...], "params": {...}},
+            {"match": {"contains": "refund"}, "params": {"temperature": 0.0}},
+        ]):
+            run_pipeline()
+    """
+
+    def __init__(self, override_specs: Optional[List[Dict[str, Any]]] | Dict[str, Any] = None):
+        if isinstance(override_specs, dict):
+            override_specs = [override_specs]
+        self._specs = override_specs or []
+        self._tok = None
+
+    def __enter__(self):
+        self._tok = set_override_specs(self._specs)
+        return self
+
+    def __exit__(self, exc_type, exc, tb):
+        if self._tok is not None:
+            clear_override_specs(self._tok)