safetynet/apps/inspector/semantic.py

"""Pinned offline EmbeddingGemma 2 GGUF worker. No downloads or HTTP inference."""
import atexit
import base64
import hashlib
import io
import json
import math
import os
from pathlib import Path
import queue
import re
import subprocess
import threading
import time
import wave

MODEL_ID="google/embeddinggemma-2"
REVISION="bfcd298762cc34d0357ece5ebdd31791a3a374d8"
DIMENSIONS=768
PROFILE="document-v1"
RUNTIME_REVISION="b11461"
RUNTIME_BUILD="b11461-no-logits-v1"
MAX_CHUNKS=64
MAX_LINE=4*1024*1024
MAX_REQUEST=16*1024*1024
MODEL_FILES={
    "embeddinggemma-2-Q8_0.gguf":(309855456,"2188ac1deca4b77dffefd603c2776a9d76d9d74ec01841392982ebb840b09135"),
    "mmproj-embeddinggemma-2-Q8_0.gguf":(554821024,"c4a8a52691ecef40618438928bdf9e68379b854e24166f292592353db0aab64f")}


def verify_file(directory,name,size,digest):
    if not isinstance(name,str) or not name or Path(name).is_absolute() or ".." in Path(name).parts: raise ValueError("Invalid local model path")
    file=(directory/name).resolve()
    if not file.is_relative_to(directory) or not file.is_file() or file.stat().st_size!=size: raise ValueError("Local semantic artifact unavailable")
    if not re.fullmatch(r"[0-9a-f]{64}",digest): raise ValueError("Invalid artifact digest")
    calculated=hashlib.sha256()
    with file.open("rb") as source:
        while part:=source.read(1024*1024): calculated.update(part)
    if calculated.hexdigest()!=digest: raise ValueError("Local semantic artifact verification failed")
    return file


def validate_vectors(vectors,expected=None):
    if not isinstance(vectors,list) or not 1<=len(vectors)<=MAX_CHUNKS or expected is not None and len(vectors)!=expected: raise ValueError("Invalid semantic vector batch")
    for vector in vectors:
        if not isinstance(vector,list) or len(vector)!=DIMENSIONS or any(not isinstance(value,(int,float)) or isinstance(value,bool) or not math.isfinite(value) for value in vector): raise ValueError("Invalid semantic output")
        norm=math.sqrt(sum(value*value for value in vector))
        if not 0.99<=norm<=1.01: raise ValueError("Semantic output is not normalized")
    return vectors


def validate_readiness(ready):
    if not isinstance(ready,dict) or ready.get("ready") is not True or ready.get("model")!=MODEL_ID or ready.get("dimensions")!=DIMENSIONS or ready.get("runtime")!=RUNTIME_BUILD or ready.get("device")!="cpu" or ready.get("vulkan_module_loaded") is not False:
        raise ValueError("Local semantic worker identity mismatch")


class NativeSemanticModel:
    def __init__(self,directory,worker_path):
        self.directory=Path(directory).resolve(); self.worker_path=Path(worker_path).resolve()
        manifest_path=self.directory/"model-manifest.json"
        if not manifest_path.is_file() or manifest_path.stat().st_size>65536: raise ValueError("Local model manifest unavailable")
        manifest=json.loads(manifest_path.read_text(encoding="utf-8"))
        if manifest.get("model_id")!=MODEL_ID or manifest.get("revision")!=REVISION or manifest.get("dimensions")!=DIMENSIONS or manifest.get("quantization")!="Q8_0" or manifest.get("runtime_tag")!=RUNTIME_REVISION: raise ValueError("Pinned Google model identity mismatch")
        self.model=verify_file(self.directory,*next((name,*value) for name,value in MODEL_FILES.items() if not name.startswith("mmproj")))
        self.projector=verify_file(self.directory,*next((name,*value) for name,value in MODEL_FILES.items() if name.startswith("mmproj")))
        runtime_path=self.worker_path.parent/"runtime-manifest.json"
        if not runtime_path.is_file() or runtime_path.stat().st_size>65536: raise ValueError("Verified native runtime manifest unavailable")
        runtime=json.loads(runtime_path.read_text(encoding="utf-8")); files=runtime.get("files")
        if not isinstance(files,list) or not 1<=len(files)<=64 or not any(item.get("path")==self.worker_path.name for item in files): raise ValueError("Invalid native runtime manifest")
        for item in files: verify_file(self.worker_path.parent,item["path"],item["bytes"],item["sha256"])
        # All DLLs load from the private, verified runtime directory.
        verified={item["path"].casefold() for item in files}
        if any(file.name.casefold() not in verified for file in self.worker_path.parent.glob("*.dll")): raise ValueError("Unverified runtime module")
        self.lock=threading.Lock();self.process=None;self.responses=None
        atexit.register(self.close)

    def info(self):
        return {"model":MODEL_ID,"revision":REVISION,"dimensions":DIMENSIONS,"profile":PROFILE,"artifactFormat":"gguf","quantization":"Q8_0","runtimeRevision":RUNTIME_REVISION,"local":True,"backend":"native-gguf","device":"cpu","modalities":["text","image","audio"]}

    def close(self):
        process=self.process;self.process=None
        if process is not None:
            try: process.kill();process.wait(timeout=3)
            except Exception: pass
            for stream in (process.stdin,process.stdout):
                try: stream.close()
                except Exception: pass

    def _read(self,process,responses):
        try:
            while True:
                line=process.stdout.readline(MAX_LINE+1)
                if not line or len(line)>MAX_LINE or not line.endswith(b"\n"): raise ValueError()
                response=json.loads(line)
                responses.put(response,timeout=1)
        except Exception:
            try: responses.put({"protocolFailure":True},timeout=1)
            except Exception: pass

    def _receive(self,timeout):
        try:
            response=self.responses.get(timeout=timeout)
            if not isinstance(response,dict) or response.get("protocolFailure"): raise ValueError()
            return response
        except Exception:
            self.close();raise ValueError("Local semantic worker failed or timed out") from None

    def _start(self):
        if self.process is not None and self.process.poll() is None: return
        self.close();self.responses=queue.Queue(maxsize=2)
        environment={key:value for key,value in os.environ.items() if key.upper() in {"SYSTEMROOT","WINDIR","TEMP","TMP"}}
        environment.update({"SEARCHI_FORCE_CPU":"1","GGML_CPU_ONLY":"1","GGML_DISABLE_VULKAN":"1"})
        self.process=subprocess.Popen([str(self.worker_path),"--model",str(self.model),"--threads","2","--device","cpu"],
                                      cwd=str(self.worker_path.parent),stdin=subprocess.PIPE,stdout=subprocess.PIPE,stderr=subprocess.DEVNULL,
                                      env=environment,creationflags=getattr(subprocess,"CREATE_NO_WINDOW",0))
        threading.Thread(target=self._read,args=(self.process,self.responses),daemon=True).start()
        ready=self._receive(20)
        try: validate_readiness(ready)
        except Exception: self.close();raise

    def _request(self,request,timeout=30):
        payload=json.dumps(request,ensure_ascii=False,allow_nan=False).encode("utf-8")+b"\n"
        if len(payload)>=MAX_REQUEST: raise ValueError("Semantic request budget exceeded")
        with self.lock:
            self._start()
            process=self.process
            # A blocked pipe write is also bounded and tears down the process.
            timer=threading.Timer(timeout,lambda:process.kill() if process.poll() is None else None);timer.daemon=True;timer.start()
            try:
                process.stdin.write(payload);process.stdin.flush();response=self._receive(timeout)
                if response.get("ok") is not True or response.get("device")!="cpu": raise ValueError("Local semantic inference failed")
                return response
            except Exception:
                self.close();raise ValueError("Local semantic inference failed") from None
            finally: timer.cancel();payload=None

    def text_vectors(self,text):
        if not isinstance(text,str) or not text.strip() or len(text.encode("utf-8"))>1024*1024: raise ValueError("Semantic text budget exceeded")
        encoded=text.encode("utf-8");windows=self._request({"op":"windows","text":text,"title":"none"}).get("windows")
        if not isinstance(windows,list) or not 1<=len(windows)<=MAX_CHUNKS: raise ValueError("Semantic window budget exceeded")
        previous_end=0;parts=[]
        for item in windows:
            start=item.get("start_byte");end=item.get("end_byte");part=item.get("text")
            if not isinstance(start,int) or not isinstance(end,int) or not isinstance(part,str) or start<0 or start>previous_end or end<=previous_end or end>len(encoded) or part.encode("utf-8")!=encoded[start:end]: raise ValueError("Semantic source coverage invalid")
            previous_end=end;parts.append(part)
        if previous_end!=len(encoded): raise ValueError("Semantic source tail missing")
        vectors=[];deadline=time.monotonic()+60
        for index in range(0,len(parts),8):
            remaining=deadline-time.monotonic()
            if remaining<=0: raise ValueError("Semantic inference budget exceeded")
            batch=parts[index:index+8];response=self._request({"op":"documents","texts":batch,"title":"none"},min(30,remaining))
            vectors.extend(validate_vectors(response.get("vectors"),len(batch)))
        return vectors

    def image_vectors(self,data):
        from PIL import Image,ImageSequence,ImageOps
        image=Image.open(io.BytesIO(data));frames=getattr(image,"n_frames",1)
        if frames>32 or image.width*image.height*frames>50_000_000: raise ValueError("Semantic image budget exceeded")
        vectors=[]
        for frame in ImageSequence.Iterator(image):
            frame=ImageOps.exif_transpose(frame).convert("RGBA")
            background=Image.new("RGBA",frame.size,"white");background.alpha_composite(frame);frame=background.convert("RGB")
            scale=min(1,672/max(frame.size));width=max(48,round(frame.width*scale/48)*48);height=max(48,round(frame.height*scale/48)*48)
            frame=frame.resize((width,height),Image.Resampling.LANCZOS)
            request={"op":"image","projector":str(self.projector),"width":width,"height":height,"data":base64.b64encode(frame.tobytes()).decode("ascii")}
            vectors.extend(validate_vectors(self._request(request).get("vectors"),1))
        return vectors

    def audio_vectors(self,data):
        import numpy as np
        with wave.open(io.BytesIO(data),"rb") as source:
            if source.getnchannels()!=1 or source.getframerate()!=16000 or source.getsampwidth()!=2 or not 0<source.getnframes()<=480000: raise ValueError("Only bounded mono16kHz PCM16 WAV is supported")
            payload=source.readframes(source.getnframes())
            if len(payload)!=source.getnframes()*2: raise ValueError("Truncated audio data")
            pcm=np.frombuffer(payload,dtype="<i2").astype("<f4")/32768
        request={"op":"audio","projector":str(self.projector),"data":base64.b64encode(pcm.tobytes()).decode("ascii")}
        return validate_vectors(self._request(request).get("vectors"),1)