Artifex/control_plane/model_studio/profiles.py

from __future__ import annotations

import hashlib
import json
import shlex
import subprocess
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Protocol


@dataclass
class ArchaeologyFinding:
    subject: str
    confidence: str
    evidence: list[str]
    details: dict[str, Any]


class ModelProjectProfile(Protocol):
    name: str
    def archaeology(self, repository_path: str) -> dict[str, Any]: ...
    def validate_evaluation_suite(self, repository_path: str, suite_reference: str) -> dict[str, Any]: ...
    def training_command(self, recipe: dict[str, Any], output_directory: str) -> list[str]: ...


class GuardModelProfile:
    name = "guard"
    base_model = "Qwen/Qwen2.5-Coder-3B-Instruct"

    def archaeology(self, repository_path: str) -> dict[str, Any]:
        root = Path(repository_path)
        findings: list[dict[str, Any]] = []
        def add(subject: str, confidence: str, evidence: list[Path], **details: Any) -> None:
            findings.append({"subject": subject, "confidence": confidence, "evidence": [str(item) for item in evidence], "details": details})

        trainer = root / "scripts" / "05_finetune.py"
        evaluator = root / "scripts" / "run_evmbench.py"
        strict_scorer = root / "scripts" / "score_evmbench_strict.py"
        add("base_model", "CONFIRMED" if trainer.exists() else "UNKNOWN", [trainer] if trainer.exists() else [], value=self.base_model)
        add("training_entrypoint", "CONFIRMED" if trainer.exists() else "UNKNOWN", [trainer] if trainer.exists() else [], command="python scripts/05_finetune.py --model ... --training-data ... --output ...")
        add("evaluation_entrypoint", "CONFIRMED" if evaluator.exists() else "UNKNOWN", [item for item in [evaluator, strict_scorer] if item.exists()], command="python scripts/run_evmbench.py ...", strict_scoring=bool(strict_scorer.exists()))
        checkpoints = []
        for adapter in sorted((root / "models").glob("**/adapter_model.*")) if (root / "models").exists() else []:
            checkpoints.append({"name": adapter.parent.name, "reference": str(adapter.parent), "hash": self._hash_file(adapter), "adapter_type": "LORA", "confidence": "CONFIRMED"})
        datasets = []
        for manifest in sorted((root / "data" / "production_authorized").glob("*.json")) if (root / "data" / "production_authorized").exists() else []:
            try:
                payload = json.loads(manifest.read_text(encoding="utf-8"))
                count = len(payload) if isinstance(payload, list) else payload.get("record_count")
            except (OSError, json.JSONDecodeError):
                count = None
            datasets.append({"name": manifest.stem, "reference": str(manifest), "hash": self._hash_file(manifest), "record_count": count, "tags": ["imported", "training"]})
        reports = [str(path) for path in sorted((root / "results").glob("qwen25_coder_3b_*") if (root / "results").exists() else [])]
        benchmark = root / "evmbench"
        add("benchmark_checkout", "CONFIRMED" if benchmark.exists() else "UNKNOWN", [benchmark] if benchmark.exists() else [], strict_scorer_exists=strict_scorer.exists())
        return {"root_exists": root.exists(), "findings": findings, "checkpoints": checkpoints, "datasets": datasets, "historical_reports": reports, "benchmark_reference": str(benchmark), "trainer_reference": str(trainer), "evaluator_reference": str(evaluator), "strict_scorer_reference": str(strict_scorer)}

    def validate_evaluation_suite(self, repository_path: str, suite_reference: str) -> dict[str, Any]:
        root = Path(repository_path)
        evaluator = root / "scripts" / "run_evmbench.py"
        scorer = root / "scripts" / "score_evmbench_strict.py"
        benchmark = Path(suite_reference)
        available = evaluator.exists() and scorer.exists() and benchmark.exists()
        return {"status": "WARNING" if available else "BLOCKED", "evidence": {"evaluator": str(evaluator), "strict_scorer": str(scorer), "benchmark": str(benchmark), "reason": "Scripts and checkout are present, but no fresh Spark baseline has yet proven model loading, output parsing, reproducibility, or contamination checks." if available else "Guard evaluation harness or strict scorer missing."}}

    def training_command(self, recipe: dict[str, Any], output_directory: str) -> list[str]:
        dataset = str(recipe["training_data"])
        return ["python", "scripts/05_finetune.py", "--model", str(recipe.get("base_model", self.base_model)), "--training-data", dataset, "--output", output_directory]

    @staticmethod
    def _hash_file(path: Path) -> str:
        digest = hashlib.sha256()
        with path.open("rb") as handle:
            for block in iter(lambda: handle.read(1024 * 1024), b""):
                digest.update(block)
        return digest.hexdigest()


class SparkGuardDatasetInventory:
    """Read-only manifest metadata probe for a registered Spark Guard workspace."""

    def __init__(self, ssh_alias: str = "spark") -> None:
        self.ssh_alias = ssh_alias

    def inspect(self, references: list[str]) -> list[dict[str, Any]]:
        script = """import hashlib,json,os,sys
keys=['dataset_status','final_model_eligible','production_training_authorized','benchmark_source_included','synthetic_code_included','split','c4_invalid','teacher','vulnerability_type']
for path in sys.argv[1:]:
 try:
  raw=open(path,'rb').read(); value=json.loads(raw); first=value[0] if isinstance(value,list) and value else {}
  print(json.dumps({'reference':path,'content_hash':hashlib.sha256(raw).hexdigest(),'record_count':len(value) if isinstance(value,list) else None,'first':{key:first.get(key) for key in keys if key in first}},ensure_ascii=True))
 except Exception as exc: print(json.dumps({'reference':path,'error':str(exc)},ensure_ascii=True))
"""
        if not references:
            return []
        output = self._remote("python3 -c " + shlex.quote(script) + " " + " ".join(self._quote(item) for item in references), timeout=120)
        records = []
        for line in output.splitlines():
            try:
                item = json.loads(line)
            except json.JSONDecodeError:
                continue
            if "content_hash" in item and item.get("record_count") is not None:
                records.append(item)
        return records

    def _remote(self, command: str, *, timeout: int = 120) -> str:
        completed = subprocess.run(["ssh", self.ssh_alias, command], capture_output=True, text=True, encoding="utf-8", errors="replace", timeout=timeout, check=False)
        if completed.returncode:
            raise RuntimeError(completed.stderr.strip() or f"Spark inventory command failed: {command}")
        return completed.stdout.strip()

    @staticmethod
    def _quote(value: str) -> str:
        return "'" + value.replace("'", "'\\''") + "'"


class SparkGuardDatasetMaterializer(SparkGuardDatasetInventory):
    """Creates new versioned JSON manifests from authorized Guard source records on Spark."""

    def materialize(self, source_references: list[str], output_directory: str, *, strict: bool = False) -> dict[str, Any]:
        script = """import hashlib,json,os,sys
out=sys.argv[1]; strict=sys.argv[2]=='1'; sources=sys.argv[3:]; os.makedirs(out,exist_ok=True)
train=[]; validation=[]; regression=[]; seen=set(); stats={'sources':len(sources),'read':0,'accepted':0,'malformed':0,'duplicates':0}
for source in sources:
 value=json.load(open(source,encoding='utf-8'))
 if not isinstance(value,list): continue
 for record in value:
  stats['read']+=1
  if not isinstance(record,dict) or not isinstance(record.get('input'),str) or not record['input'].strip() or not isinstance(record.get('output'),str) or not record['output'].strip(): stats['malformed']+=1; continue
  if strict:
   try: target=json.loads(record['output'])
   except Exception: stats['malformed']+=1; continue
   if not isinstance(target,dict) or not isinstance(target.get('findings'),list) or record.get('c4_invalid') is True or 'pragma solidity' not in record['input'].lower(): stats['malformed']+=1; continue
   normalized=[]
   for finding in target['findings']:
    if not isinstance(finding,dict): continue
    kind=finding.get('canonical_type') or finding.get('type')
    if kind: normalized.append({key:finding[key] for key in ['canonical_type','type','severity','location','description','evidence','vulnerable_line_spans'] if key in finding})
   record={**record,'output':json.dumps({'findings':normalized},ensure_ascii=False),'artifex_schema':'guard_finding_v1'}
  source_key=str(record.get('source_sha256') or record.get('parent_source_sha256') or hashlib.sha256(record['input'].encode()).hexdigest())
  key=hashlib.sha256((record['input']+'\\0'+record['output']).encode()).hexdigest()
  if key in seen: stats['duplicates']+=1; continue
  seen.add(key); record={**record,'artifex_source_key':source_key,'artifex_record_hash':key,'artifex_source_manifest':source}; bucket=int(hashlib.sha256(source_key.encode()).hexdigest(),16)%10
  split='train' if bucket<8 else 'validation' if bucket<9 else 'regression'; record['split']=split; (train if split=='train' else validation if split=='validation' else regression).append(record); stats['accepted']+=1
for name,rows in [('train',train),('validation',validation),('regression',regression),('training',train+validation)]:
 path=os.path.join(out,name+'.json'); json.dump(rows,open(path,'w',encoding='utf-8'),ensure_ascii=False); stats[name]=len(rows); stats[name+'_sha256']=hashlib.sha256(open(path,'rb').read()).hexdigest()
json.dump(stats,open(os.path.join(out,'manifest.json'),'w',encoding='utf-8'),indent=2); print(json.dumps({'output_directory':out,**stats}))
"""
        if not source_references:
            raise ValueError("Dataset materialization requires source manifests.")
        output = self._remote("python3 -c " + shlex.quote(script) + " " + self._quote(output_directory) + " " + ("1" if strict else "0") + " " + " ".join(self._quote(item) for item in source_references), timeout=900)
        return json.loads(output)

    def sample(self, reference: str, *, count: int = 8, maximum_field_chars: int = 3000) -> list[dict[str, Any]]:
        script = """import json,sys
rows=json.load(open(sys.argv[1],encoding='utf-8')); count=int(sys.argv[2]); limit=int(sys.argv[3]); step=max(1,len(rows)//max(1,count)); out=[]
for index in range(0,len(rows),step):
 record=rows[index]; out.append({key:(value[:limit] if isinstance(value,str) else value) for key,value in record.items() if key in ['input','output','record_id','source_sha256','parent_source_sha256','vulnerability_type','artifex_source_manifest']})
 if len(out)>=count: break
print(json.dumps(out,ensure_ascii=True))
"""
        output = self._remote("python3 -c " + shlex.quote(script) + " " + self._quote(reference) + f" {count} {maximum_field_chars}", timeout=120)
        return json.loads(output)