Add targeting

StarDylan · Nov 17, 2023 · 48d1ad2 · 48d1ad2
1 parent cfe4d42
commit 48d1ad2
Show file tree

Hide file tree

Showing 2 changed files with 73 additions and 1 deletion.
diff --git a/requirements.txt b/requirements.txt
@@ -5,4 +5,5 @@ sqlalchemy==2.0.7
 psycopg2-binary~=2.9.3
 python-dotenv
 pre-commit
-pycryptodome
+pycryptodome
+requests
diff --git a/src/api/songs.py b/src/api/songs.py
@@ -2,6 +2,9 @@
 from pydantic import BaseModel
 import sqlalchemy
 from src import database as db
+import json
+import requests
+import os
 
 from src.api.models import SongPlayLink
 
@@ -124,6 +127,70 @@ def remove_song(song_id: int, authorization: SongAuthorization):
     return "Song Removed"
 
 
+def play_ad_if_needed(conn, user_id) -> str | None:
+    result = conn.execute(sqlalchemy.text("""
+            SELECT COUNT(*) FROM song_history
+            WHERE user_id = :user_id
+            """),
+            [{
+                "user_id": user_id
+            }]).scalar_one()
+
+    if result < 5:
+        return None
+
+    result = conn.execute(sqlalchemy.text("""
+            SELECT song_name, artist
+            FROM song_history
+            JOIN songs ON song_history.song_id = songs.id
+            WHERE user_id = 23
+            ORDER BY song_history.created_at DESC
+            LIMIT 5
+            """)).all()
+
+    song_prompt = "Songs:\n"
+    for song in result:
+        song_prompt += song.song_name + " by " + song.artist + "\n"
+
+    payload = json.dumps({
+        "model": "llama2-uncensored",
+        "system": "Classify the user's mood based on the following song titles into only one of these emotions: Happy, Sad, Angry. Only include the classification as one word.",
+        "prompt": song_prompt,
+        "stream": False
+        })
+    headers = {
+    'Content-Type': 'application/json'
+    }
+
+    print("Getting Sentiment from OLLAMA")
+
+    response = requests.request("POST", os.environ.get("OLLAMA_URI"), headers=headers, data=payload)
+    response = response.json()
+    print(response)
+
+    mood = ""
+    if "happy" in response["response"].lower():
+        mood = "HAPPY"
+    elif "sad" in response["response"].lower():
+        mood = "SAD"
+    elif "angry" in response["response"].lower():
+        mood = "ANGRY"
+
+    if mood == "":
+        return None
+
+    result = conn.execute(sqlalchemy.text("""
+            SELECT link FROM ad_campaigns
+            WHERE target_mood = :mood
+            ORDER BY RANDOM() """), [{
+                "mood": mood
+            }]).scalar_one_or_none()
+
+    if result is None:
+        return None
+
+    return result
+
 @router.get("/{song_id}/play")
 def play_song(song_id: int, user_id: str = Header(None)) -> SongPlayLink:
     """ """
@@ -168,6 +235,10 @@ def play_song(song_id: int, user_id: str = Header(None)) -> SongPlayLink:
 
             return "Song not available on user's platform"
 
+        ad_link = play_ad_if_needed(conn, user_id)
+        if ad_link is not None:
+            return ad_link
+
         conn.execute(sqlalchemy.text("""
           INSERT INTO song_history (user_id, song_id) VALUES (:user_id, :song_id)
             """),