102 lines
3.1 KiB
Python
102 lines
3.1 KiB
Python
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()
|