Update database.py

database.py

@@ -2,17 +2,28 @@
 import sqlite3
 import datetime
 import uuid
-from typing import List, Dict, Any
+from typing import List, Dict, Any, Optional
 
+# --- Configuration ---
 DB_FILE = "social_media_platform.db"
 
-# A simple in-memory cache for agent data to avoid constant DB lookups
+# A simple in-memory cache to store agent data (token -> {id, name})
+# This prevents hitting the database for every single API call.
 AGENT_CACHE: Dict[str, Any] = {}
 
+# ===================================================================
+#                      Initialization and Setup
+# ===================================================================
+
 def initialize_db():
-    conn = sqlite3.connect(DB_FILE)
+    """
+    Initializes the database by creating all necessary tables if they don't already exist.
+    This function is safe to call every time the server starts.
+    """
+    conn = sqlite3.connect(DB_FILE, check_same_thread=False)
     cursor = conn.cursor()
-    # Create agents table with a secret API token
+    
+    # agents table stores the permanent state of each AI agent, including their secret token for API access.
     cursor.execute("""
     CREATE TABLE IF NOT EXISTS agents (
         agent_id INTEGER PRIMARY KEY,
@@ -22,31 +33,65 @@ def initialize_db():
         api_token TEXT NOT NULL UNIQUE
     )
     """)
-    # Other tables
-    cursor.execute("CREATE TABLE IF NOT EXISTS posts (post_id INTEGER PRIMARY KEY, agent_id INTEGER, content TEXT, timestamp TEXT)")
-    cursor.execute("CREATE TABLE IF NOT EXISTS comments (comment_id INTEGER PRIMARY KEY, post_id INTEGER, agent_id INTEGER, content TEXT, timestamp TEXT)")
-    cursor.execute("CREATE TABLE IF NOT EXISTS likes (like_id INTEGER PRIMARY KEY, post_id INTEGER, agent_id INTEGER)")
+    
+    # posts table stores the core content.
+    cursor.execute("""
+    CREATE TABLE IF NOT EXISTS posts (
+        post_id INTEGER PRIMARY KEY,
+        agent_id INTEGER NOT NULL,
+        content TEXT NOT NULL,
+        timestamp TEXT NOT NULL
+    )
+    """)
+
+    # comments table links to the posts table via post_id.
+    cursor.execute("""
+    CREATE TABLE IF NOT EXISTS comments (
+        comment_id INTEGER PRIMARY KEY,
+        post_id INTEGER NOT NULL,
+        agent_id INTEGER NOT NULL,
+        content TEXT NOT NULL,
+        timestamp TEXT NOT NULL,
+        FOREIGN KEY (post_id) REFERENCES posts (post_id)
+    )
+    """)
+
+    # likes table is a simple mapping of which agent liked which post.
+    cursor.execute("""
+    CREATE TABLE IF NOT EXISTS likes (
+        like_id INTEGER PRIMARY KEY,
+        post_id INTEGER NOT NULL,
+        agent_id INTEGER NOT NULL,
+        FOREIGN KEY (post_id) REFERENCES posts (post_id),
+        UNIQUE(post_id, agent_id) -- Ensures an agent can only like a post once
+    )
+    """)
+    
     conn.commit()
     conn.close()
 
 def populate_initial_agents():
-    """Adds the 10 starting agents if they don't exist and prints their tokens."""
+    """
+    Adds a predefined list of 10 agents to the database if they aren't already there.
+    Generates a unique API token for each one and prints them to the console.
+    This is crucial for you to be able to test the API.
+    """
     agents_to_create = [
         {"name": "TechieTom", "persona": "A cynical but brilliant software developer.", "interests": "AI,programming"},
         {"name": "NatureNina", "persona": "A cheerful nature photographer.", "interests": "hiking,animals"},
         {"name": "ChefCarlo", "persona": "A passionate Italian chef.", "interests": "pasta,wine"},
         {"name": "AstroAlex", "persona": "An enthusiastic astronomer.", "interests": "stars,planets"},
         {"name": "HistoryHank", "persona": "A meticulous historian.", "interests": "battles,empires"},
-        {"name": "Fit Fiona", "persona": "A motivational fitness coach.", "interests": "running,yoga"},
+        {"name": "FitFiona", "persona": "A motivational fitness coach.", "interests": "running,yoga"},
         {"name": "GamerGabe", "persona": "A competitive esports player.", "interests": "rpgs,strategy"},
         {"name": "ArtistAnna", "persona": "A thoughtful abstract painter.", "interests": "color theory,canvases"},
         {"name": "MusoMike", "persona": "A laid-back indie musician.", "interests": "guitar,vinyl"},
         {"name": "BookwormBella", "persona": "An avid reader of classic literature.", "interests": "novels,poetry"},
     ]
     
-    conn = sqlite3.connect(DB_FILE)
+    conn = sqlite3.connect(DB_FILE, check_same_thread=False)
     cursor = conn.cursor()
-    print("\n--- Agent API Tokens ---")
+    print("\n--- Agent API Tokens (save these!) ---")
     for agent_data in agents_to_create:
         cursor.execute("SELECT api_token FROM agents WHERE name = ?", (agent_data["name"],))
         result = cursor.fetchone()
@@ -56,33 +101,151 @@ def populate_initial_agents():
                 "INSERT INTO agents (name, persona, interests, api_token) VALUES (?, ?, ?, ?)",
                 (agent_data["name"], agent_data["persona"], agent_data["interests"], token)
             )
-            print(f"  {agent_data['name']}: {token}")
+            print(f'  "{agent_data["name"]}": "{token}"')
         else:
-            print(f"  {agent_data['name']}: {result[0]} (already exists)")
+            print(f'  "{agent_data["name"]}": "{result[0]}" (already exists)')
     conn.commit()
     conn.close()
-    print("------------------------\n")
+    print("--------------------------------------\n")
 
-def get_agent_by_token(token: str) -> Dict[str, Any] | None:
+# ===================================================================
+#                           Core Data Access Functions
+# ===================================================================
+
+def get_agent_by_token(token: str) -> Optional[Dict[str, Any]]:
+    """
+    Finds an agent by their API token, using a cache to improve performance.
+    Returns a dictionary with agent_id and name, or None if not found.
+    """
     if token in AGENT_CACHE:
         return AGENT_CACHE[token]
-    conn = sqlite3.connect(DB_FILE)
+    
+    conn = sqlite3.connect(DB_FILE, check_same_thread=False)
     cursor = conn.cursor()
     cursor.execute("SELECT agent_id, name FROM agents WHERE api_token = ?", (token,))
-    agent = cursor.fetchone()
+    agent_row = cursor.fetchone()
     conn.close()
-    if agent:
-        agent_data = {"agent_id": agent[0], "name": agent[1]}
+    
+    if agent_row:
+        agent_data = {"agent_id": agent_row[0], "name": agent_row[1]}
         AGENT_CACHE[token] = agent_data
         return agent_data
     return None
 
+def _get_agent_name(agent_id: int, cursor) -> str:
+    """A helper function to get an agent's name from their ID. Handles the special HumanUser case."""
+    if agent_id == 0:
+        return "HumanUser"
+    cursor.execute("SELECT name FROM agents WHERE agent_id = ?", (agent_id,))
+    result = cursor.fetchone()
+    return result[0] if result else "Unknown Agent"
+
 def get_timeline(limit: int = 20) -> List[Dict[str, Any]]:
-    conn = sqlite3.connect(DB_FILE)
-    # This makes the output rows behave like dictionaries
-    conn.row_factory = sqlite3.Row
+    """
+    Fetches a list of recent posts for the main timeline.
+    For each post, it calculates the number of likes and comments using subqueries.
+    """
+    conn = sqlite3.connect(DB_FILE, check_same_thread=False)
+    conn.row_factory = sqlite3.Row  # This allows accessing columns by name
     cursor = conn.cursor()
+    
     query = """
     SELECT
-        p.post_id, p.content, p.timestamp,
-        a.agent_id, a
+        p.post_id, p.content, p.timestamp, p.agent_id,
+        a.name AS author_name,
+        (SELECT COUNT(*) FROM likes WHERE post_id = p.post_id) AS likes_count,
+        (SELECT COUNT(*) FROM comments WHERE post_id = p.post_id) AS comments_count
+    FROM posts p
+    LEFT JOIN agents a ON p.agent_id = a.agent_id
+    ORDER BY p.timestamp DESC
+    LIMIT ?
+    """
+    
+    posts_data = []
+    for row in cursor.execute(query, (limit,)):
+        # Handle the case where the author is the Human (agent_id=0), who isn't in the agents table
+        author_name = row['author_name'] if row['author_name'] else "HumanUser"
+        posts_data.append({
+            "post_id": row["post_id"],
+            "author": {"agent_id": row["agent_id"], "name": author_name},
+            "content": row["content"],
+            "timestamp": datetime.datetime.fromisoformat(row["timestamp"]),
+            "stats": {"likes": row["likes_count"], "comments": row["comments_count"]}
+        })
+    conn.close()
+    return posts_data
+
+def create_post(agent_id: int, content: str, agent_name: str) -> Dict[str, Any]:
+    """
+    Adds a new post to the database.
+    Returns a dictionary representing the newly created post.
+    """
+    conn = sqlite3.connect(DB_FILE, check_same_thread=False)
+    cursor = conn.cursor()
+    
+    timestamp = datetime.datetime.utcnow().isoformat()
+    cursor.execute(
+        "INSERT INTO posts (agent_id, content, timestamp) VALUES (?, ?, ?)",
+        (agent_id, content, timestamp)
+    )
+    post_id = cursor.lastrowid
+    conn.commit()
+    conn.close()
+    
+    return {
+        "post_id": post_id,
+        "author": {"agent_id": agent_id, "name": agent_name},
+        "content": content,
+        "timestamp": datetime.datetime.fromisoformat(timestamp),
+        "stats": {"likes": 0, "comments": 0}
+    }
+
+def create_comment(post_id: int, agent_id: int, content: str, agent_name: str) -> Optional[Dict[str, Any]]:
+    """
+    Adds a new comment to a specific post.
+    Returns a dictionary of the new comment, or None if the post doesn't exist.
+    """
+    conn = sqlite3.connect(DB_FILE, check_same_thread=False)
+    cursor = conn.cursor()
+
+    # First, check if the post exists
+    cursor.execute("SELECT post_id FROM posts WHERE post_id = ?", (post_id,))
+    if cursor.fetchone() is None:
+        conn.close()
+        return None
+
+    timestamp = datetime.datetime.utcnow().isoformat()
+    cursor.execute(
+        "INSERT INTO comments (post_id, agent_id, content, timestamp) VALUES (?, ?, ?, ?)",
+        (post_id, agent_id, content, timestamp)
+    )
+    comment_id = cursor.lastrowid
+    conn.commit()
+    conn.close()
+
+    return {
+        "comment_id": comment_id,
+        "author": {"agent_id": agent_id, "name": agent_name},
+        "content": content,
+        "timestamp": datetime.datetime.fromisoformat(timestamp)
+    }
+
+def create_like(post_id: int, agent_id: int) -> bool:
+    """
+    Adds a 'like' to a specific post for a given agent.
+    Returns True if the like was added, False if it already existed or post not found.
+    """
+    conn = sqlite3.connect(DB_FILE, check_same_thread=False)
+    cursor = conn.cursor()
+    try:
+        # The UNIQUE constraint on (post_id, agent_id) will cause an IntegrityError if a duplicate is inserted.
+        cursor.execute("INSERT INTO likes (post_id, agent_id) VALUES (?, ?)", (post_id, agent_id))
+        conn.commit()
+        success = True
+    except sqlite3.IntegrityError:
+        # This means the like already existed, or the post_id is invalid. We can treat it as a "success"
+        # from the client's perspective, as the end state is that the post is liked.
+        success = True 
+    finally:
+        conn.close()
+    return success