Update app.py

app.py

@@ -1,4 +1,6 @@
 import os
+import threading
+import time
 import datetime
 from fastapi import FastAPI
 from supabase import create_client, Client
@@ -6,44 +8,68 @@ from transformers import pipeline
 from huggingface_hub import login
 import torch
 
+# --- CONFIGURATION ---
+AGENT_NAME = "Nemotron-Validator-Pod"
+MODEL_ID = "nvidia/Llama-3.1-Minitron-4B-Width-Base"
+
 app = FastAPI()
-url = os.environ.get("SUPABASE_URL")
-key = os.environ.get("SUPABASE_KEY")
-hf_token = os.environ.get("HF_TOKEN")
-supabase: Client = create_client(url, key)
+supabase: Client = create_client(os.environ.get("SUPABASE_URL"), os.environ.get("SUPABASE_KEY"))
 
-if hf_token:
-    login(token=hf_token)
+if os.environ.get("HF_TOKEN"):
+    login(token=os.environ.get("HF_TOKEN"))
 
-# Load NVIDIA Nemotron-Mini-4B
+print(f"📦 Loading {MODEL_ID}...")
 pipe = pipeline(
     "text-generation",
-    model="nvidia/Llama-3.1-Minitron-4B-Width-Base", # Or nvidia/Llama-3.1-Minitron-4B-Width-Base
+    model=MODEL_ID,
     model_kwargs={"torch_dtype": torch.bfloat16},
-    device_map="auto",
-    token=hf_token
+    device_map="auto"
 )
 
-AGENT_NAME = "Nemotron-Validator-Pod"
+def worker_loop():
+    print(f"⚖️ {AGENT_NAME} Validator Loop Started.")
+    while True:
+        try:
+            # 1. Pull tasks assigned to ME for validation
+            res = supabase.table("tasks").select("*").eq("status", "processing_val").eq("assigned_to_name", AGENT_NAME).execute()
+            
+            for task in res.data:
+                task_id = task['id']
+                original_prompt = task['input_text']
+                # Get the output gemma just posted
+                gemma_content = task['output_data'].get('agent_gemma', {}).get('content', '')
+
+                print(f"🔍 Validating task {task_id[:8]}...")
+
+                # 2. Validation Inference
+                val_prompt = f"Task: {original_prompt}\nResponse: {gemma_content}\nAnalyze if this is correct. Reply with VALID or INVALID and reason."
+                
+                start_time = time.time()
+                outputs = pipe(val_prompt, max_new_tokens=150)
+                val_result = outputs[0]["generated_text"]
+                latency = round(time.time() - start_time, 2)
+
+                # 3. Update Supabase: Merge validation into output_data and flip status
+                new_output_data = task['output_data']
+                new_output_data['agent_nemotron'] = {"content": val_result, "latency": f"{latency}s"}
+
+                supabase.table("tasks").update({
+                    "status": "val_completed",
+                    "output_data": new_output_data
+                }).eq("id", task_id).execute()
+                print(f"✅ Validated task {task_id[:8]} in {latency}s")
+
+        except Exception as e:
+            print(f"⚠️ Validator Error: {e}")
+        
+        time.sleep(2)
+
+threading.Thread(target=worker_loop, daemon=True).start()
 
 @app.on_event("startup")
-def startup_event():
-    agent_data = {
-        "name": AGENT_NAME,
-        "model_name": "nvidia/Minitron-4B",
-        "status": "online",
-        "last_seen": datetime.datetime.now().isoformat(),
-        "capabilities": {"task": "validation", "specialty": "accuracy_check"}
-    }
+def register():
+    agent_data = {"name": AGENT_NAME, "status": "online", "last_seen": datetime.datetime.now().isoformat()}
     supabase.table("agents").upsert(agent_data, on_conflict="name").execute()
 
 @app.get("/")
-def health():
-    return {"status": "running", "agent": AGENT_NAME}
-
-@app.post("/validate")
-async def validate(original_prompt: str, ai_response: str):
-    # Prompt engineering to make Nemotron act as a judge
-    validation_prompt = f"Task: {original_prompt}\nResponse: {ai_response}\nIs this response accurate? Answer with YES or NO and a brief reason."
-    outputs = pipe(validation_prompt, max_new_tokens=100, do_sample=False)
-    return {"validation": outputs[0]["generated_text"]}
+def health(): return {"status": "alive", "worker": AGENT_NAME}