aa

tour-comfy/database.py

@@ -0,0 +1,101 @@
+import psycopg2
+import json
+
+DB_CONFIG = {
+    "host": "192.168.1.160",
+    "port": 5433,
+    "dbname": "dv",
+    "user": "dev",
+    "password": "dev"
+}
+
+def get_db_connection():
+    return psycopg2.connect(**DB_CONFIG)
+
+def upsert_person(filename, filepath=None, name=None, group_id=None, tags=None, embedding=None, clip_description=None):
+    conn = get_db_connection()
+    cur = conn.cursor()
+    try:
+        cur.execute("""
+            INSERT INTO person (filename, filepath, name, group_id, tags, embedding, clip_description)
+            VALUES (%s, %s, %s, %s, %s, %s, %s)
+            ON CONFLICT (filename) DO UPDATE 
+            SET filepath = COALESCE(EXCLUDED.filepath, person.filepath),
+                name = COALESCE(EXCLUDED.name, person.name),
+                group_id = COALESCE(EXCLUDED.group_id, person.group_id),
+                tags = COALESCE(EXCLUDED.tags, person.tags),
+                embedding = COALESCE(EXCLUDED.embedding, person.embedding),
+                clip_description = COALESCE(EXCLUDED.clip_description, person.clip_description);
+        """, (filename, filepath, name, group_id, json.dumps(tags) if tags else None, embedding, clip_description))
+        conn.commit()
+    finally:
+        cur.close()
+        conn.close()
+
+def get_person(filename):
+    conn = get_db_connection()
+    cur = conn.cursor()
+    try:
+        cur.execute("SELECT name, group_id, tags, embedding, clip_description, filepath FROM person WHERE filename = %s", (filename,))
+        return cur.fetchone()
+    finally:
+        cur.close()
+        conn.close()
+
+def list_persons():
+    conn = get_db_connection()
+    cur = conn.cursor()
+    try:
+        cur.execute("SELECT filename, name, group_id, clip_description FROM person")
+        return cur.fetchall()
+    finally:
+        cur.close()
+        conn.close()
+
+def search_similar(embedding, limit=10):
+    conn = get_db_connection()
+    cur = conn.cursor()
+    try:
+        # Convert embedding to string format for pgvector
+        embedding_str = "[" + ",".join(map(str, embedding)) + "]"
+        cur.execute("""
+            SELECT filename, name, group_id, clip_description, embedding <=> %s AS distance
+            FROM person
+            WHERE embedding IS NOT NULL
+            ORDER BY distance ASC
+            LIMIT %s;
+        """, (embedding_str, limit))
+        return cur.fetchall()
+    finally:
+        cur.close()
+        conn.close()
+
+def delete_person(filename):
+    conn = get_db_connection()
+    cur = conn.cursor()
+    try:
+        cur.execute("DELETE FROM person WHERE filename = %s", (filename,))
+        conn.commit()
+    finally:
+        cur.close()
+        conn.close()
+
+def delete_group(group_id):
+    conn = get_db_connection()
+    cur = conn.cursor()
+    try:
+        cur.execute("DELETE FROM person WHERE group_id = %s", (group_id,))
+        conn.commit()
+    finally:
+        cur.close()
+        conn.close()
+
+def get_group_files(group_id):
+    conn = get_db_connection()
+    cur = conn.cursor()
+    try:
+        cur.execute("SELECT filename, filepath FROM person WHERE group_id = %s", (group_id,))
+        return cur.fetchall()
+    finally:
+        cur.close()
+        conn.close()