Update app.py

app.py

@@ -1,88 +1,67 @@
-from transformers import pipeline
 from fastapi import FastAPI, Request
-import uvicorn
+import google.generativeai as genai
+import os
+import json
 
 from uagents import Agent, Context, Bureau, Model
 
-# ─── uAgent I/O MODEL ─────────────────────────────────────────────────────
+# Load Gemini API Key
+genai.configure(api_key=os.environ["GEMINI_API_KEY"])
 
+# Use Gemini Pro
+model = genai.GenerativeModel("gemini-pro")
+
+# uAgent input model
 class TextInput(Model):
     text: str
 
-# ─── LOAD EMOTION PIPELINE ─────────────────────────────────────────────────
-
-emotion_model = pipeline(
-    "text-classification",
-    model="bhadresh-savani/distilbert-base-uncased-emotion"
-)
-
-# ─── CUSTOM ANALYSIS LOGIC ─────────────────────────────────────────────────
-
+# Prompting 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 ────────────────────────────────────────────────────
-
+    prompt = f"""
+You are a mental health AI assistant. Return a JSON with:
+- self_harm
+- homicidal
+- distress
+- psychosis
+
+Scores between 0.01–0.99.
+
+Text: "{text}"
+Respond in JSON format only.
+"""
+    try:
+        response = model.generate_content(prompt)
+        raw = response.text.strip().strip("```json").strip("```")
+        return json.loads(raw)
+    except Exception as e:
+        return {
+            "self_harm": 0.01,
+            "homicidal": 0.01,
+            "distress": 0.01,
+            "psychosis": 0.01,
+            "error": str(e)
+        }
+
+# uAgent definition
 agent = Agent(name="sentiment_agent")
 
 @agent.on_message(model=TextInput)
 async def handle_message(ctx: Context, sender: str, msg: TextInput):
-    flags = analyze_text_metrics(msg.text)
-    await ctx.send(sender, flags)
-
-# ─── FASTAPI HTTP ENDPOINT ───────────────────────────────────────────────
+    result = analyze_text_metrics(msg.text)
+    await ctx.send(sender, result)
 
+# FastAPI app for Space
 app = FastAPI()
 
 @app.post("/")
-async def analyze_text(request: Request):
+async def analyze(request: Request):
     data = await request.json()
     return analyze_text_metrics(data.get("text", ""))
 
-# ─── START BOTH ───────────────────────────────────────────────────────────
-
 if __name__ == "__main__":
     bureau = Bureau()
     bureau.add(agent)
-    bureau.run_in_thread()            # serve the agent on Agentverse
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+    bureau.run_in_thread()
+
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)  # Hugging Face expects 7860