Update app.py

app.py

@@ -1,100 +1,90 @@
+from transformers import pipeline
 from fastapi import FastAPI, Request
-import google.generativeai as genai
-import os, json, re
+import uvicorn
 
 from uagents import Agent, Context, Bureau, Model
 
-# ───────────────────────────────────────────────────────────────
-#  1️⃣  GEMINI API KEY   (supports either env name)
-# ───────────────────────────────────────────────────────────────
-api_key = os.getenv("GEMINI_API_KEY") or os.getenv("Gemini")
-if not api_key:
-    raise RuntimeError(
-        "✖️ Gemini API key not found. Add it to your Space secrets "
-        "as GEMINI_API_KEY (preferred) or Gemini."
-    )
-genai.configure(api_key=api_key)
-model = genai.GenerativeModel("gemini-pro")
-
-# ───────────────────────────────────────────────────────────────
-#  2️⃣  uAgent I/O MODEL
-# ───────────────────────────────────────────────────────────────
+# ─── uAgent I/O MODEL ─────────────────────────────────────────────────────
+
 class TextInput(Model):
     text: str
 
-# ───────────────────────────────────────────────────────────────
-#  3️⃣  GEMINI-POWERED ANALYSIS
-#      – robust JSON extraction
-# ───────────────────────────────────────────────────────────────
-json_pattern = re.compile(r"\{[\s\S]*\}")
-
-def analyze_text_metrics(text: str) -> dict:
-    prompt = f"""
-You are a mental-health AI assistant.
-
-Return **ONLY** valid JSON with four keys:
-  "self_harm", "homicidal", "distress", "psychosis"
-Each value must be a number between 0.01 and 0.99 (no strings).
-
-Example:
-{{"self_harm":0.91,"homicidal":0.12,"distress":0.72,"psychosis":0.18}}
-
-Text to analyse:
-\"\"\"{text}\"\"\"
-""".strip()
-
-    try:
-        resp = model.generate_content(prompt, safety_settings={"HARASSMENT": "block_none"})
-        raw = resp.text.strip()
-
-        # grab first JSON object—even if Gemini wrapped it in ```json``` fences
-        m = json_pattern.search(raw)
-        if not m:
-            raise ValueError(f"No JSON found in Gemini response:\n{raw}")
-
-        data = json.loads(m.group(0))
-        # basic sanity check
-        for k in ("self_harm", "homicidal", "distress", "psychosis"):
-            if k not in data:
-                raise KeyError(f"Missing key {k}")
-        return data
-
-    except Exception as e:
-        # fail-safe minimal response
-        return {
-            "self_harm": 0.01,
-            "homicidal": 0.01,
-            "distress": 0.01,
-            "psychosis": 0.01,
-            "error": str(e)
-        }
-
-# ───────────────────────────────────────────────────────────────
-#  4️⃣  uAgent DEFINITION
-# ───────────────────────────────────────────────────────────────
+# ─── LOAD EMOTION PIPELINE ─────────────────────────────────────────────────
+
+emotion_model = pipeline(
+    "text-classification",
+    model="bhadresh-savani/distilbert-base-uncased-emotion"
+)
+
+# ─── CUSTOM ANALYSIS LOGIC ─────────────────────────────────────────────────
+
+def analyze_text_metrics(text):
+    results = emotion_model(text)
+    t = text.lower()
+
+    suicide_keywords   = ["kill myself", "suicidal", "die", "ending it", "pills", "overdose", "no way out"]
+    psychosis_keywords = ["voices", "hallucinate", "not real", "they’re watching me", "i’m not me"]
+
+    metrics = {
+        "self_harm": 0.0,
+        "homicidal": 0.0,
+        "distress":  0.0,
+        "psychosis": 0.0
+    }
+
+    # aggregate raw scores
+    for res in results:
+        label = res["label"]
+        score = res["score"]
+        if label == "sadness":
+            metrics["self_harm"] += score
+            metrics["distress"]  += score * 0.6
+        elif label in ("anger", "fear"):
+            metrics["homicidal"] += score
+            metrics["distress"]  += score * 0.5
+        elif label == "joy":
+            metrics["psychosis"] += score * 0.3
+        elif label == "surprise":
+            metrics["psychosis"] += score * 0.5
+
+    # keyword overrides
+    if any(kw in t for kw in suicide_keywords):
+        metrics["self_harm"] = max(metrics["self_harm"], 0.8)
+    if any(kw in t for kw in psychosis_keywords):
+        metrics["psychosis"] = max(metrics["psychosis"], 0.8)
+
+    # clamp into [0.01, 0.99]
+    for k in metrics:
+        val = metrics[k]
+        clamped = max(min(val, 0.99), 0.01)
+        metrics[k] = round(clamped, 2)
+
+    return metrics
+
+# ─── uAgent DEFINITION ────────────────────────────────────────────────────
+
 agent = Agent(name="sentiment_agent")
 
 @agent.on_message(model=TextInput)
 async def handle_message(ctx: Context, sender: str, msg: TextInput):
-    await ctx.send(sender, analyze_text_metrics(msg.text))
+    flags = analyze_text_metrics(msg.text)
+    await ctx.send(sender, flags)
+
+# ─── FASTAPI HTTP ENDPOINT ───────────────────────────────────────────────
 
-# ───────────────────────────────────────────────────────────────
-#  5️⃣  FASTAPI ENDPOINT (Hugging Face uses port 7860)
-# ───────────────────────────────────────────────────────────────
 app = FastAPI()
 
 @app.post("/")
-async def analyze(request: Request):
+async def analyze_text(request: Request):
     data = await request.json()
     return analyze_text_metrics(data.get("text", ""))
 
-# ───────────────────────────────────────────────────────────────
-#  6️⃣  START EVERYTHING
-# ───────────────────────────────────────────────────────────────
+# ─── START BOTH ───────────────────────────────────────────────────────────
+
 if __name__ == "__main__":
     bureau = Bureau()
     bureau.add(agent)
-    bureau.run_in_thread()           # expose on Agentverse
+    bureau.run_in_thread()            # serve the agent on Agentverse
+    uvicorn.run(app, host="0.0.0.0", port=8000)
+
 
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)