database.py - master - RobinDenisSh/Shvarts.D_Lab7

database.py
122 строки4 KB
RobinDenisSh
create: ai_models.py, config.py, database.py, handlers.py, main.py, ollama_client.py, prompts.db, requirements.txt
20 май 2026, 09:12
20 май 2026, 09:129fad66b
import sqlite3
import json
from datetime import datetime
from config import DB_PATH

INIT_SQL = """
CREATE TABLE IF NOT EXISTS prompts (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    name TEXT NOT NULL,
    version TEXT NOT NULL DEFAULT 'v1',
    content TEXT NOT NULL,
    target_model TEXT,
    tags TEXT,
    created_at TEXT NOT NULL,
    updated_at TEXT
);
CREATE TABLE IF NOT EXISTS versions (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    prompt_id INTEGER NOT NULL,
    version TEXT NOT NULL,
    content TEXT NOT NULL,
    created_at TEXT NOT NULL,
    FOREIGN KEY(prompt_id) REFERENCES prompts(id)
);
CREATE INDEX IF NOT EXISTS idx_prompts_name ON prompts(name);
CREATE INDEX IF NOT EXISTS idx_versions_prompt ON versions(prompt_id);
"""


def get_conn():
    conn = sqlite3.connect(DB_PATH)
    conn.row_factory = sqlite3.Row
    return conn


def init_db():
    conn = get_conn()
    conn.executescript(INIT_SQL)
    conn.commit()
    conn.close()


def add_prompt(name: str, content: str, target_model: str = "", tags: list = None) -> dict:
    conn = get_conn()
    cur = conn.cursor()
    now = datetime.utcnow().isoformat()
    tags_json = json.dumps(tags or [], ensure_ascii=False)
    # Check if prompt with this name exists to bump version
    cur.execute("SELECT version FROM prompts WHERE name = ? ORDER BY id DESC LIMIT 1", (name,))
    row = cur.fetchone()
    version = "v1"
    if row:
        try:
            prev = int(row["version"].lstrip("v"))
            version = f"v{prev + 1}"
        except Exception:
            version = "v1"
    cur.execute(
        "INSERT INTO prompts (name, version, content, target_model, tags, created_at) VALUES (?, ?, ?, ?, ?, ?)",
        (name, version, content, target_model, tags_json, now),
    )
    prompt_id = cur.lastrowid
    # Also store in versions history
    cur.execute(
        "INSERT INTO versions (prompt_id, version, content, created_at) VALUES (?, ?, ?, ?)",
        (prompt_id, version, content, now),
    )
    conn.commit()
    conn.close()
    return {"id": prompt_id, "name": name, "version": version, "content": content}


def get_prompt(prompt_id: int) -> dict:
    conn = get_conn()
    cur = conn.cursor()
    cur.execute("SELECT * FROM prompts WHERE id = ?", (prompt_id,))
    row = cur.fetchone()
    conn.close()
    if not row:
        return None
    return dict(row)


def list_prompts(name_filter: str = "") -> list:
    conn = get_conn()
    cur = conn.cursor()
    if name_filter:
        cur.execute(
            "SELECT id, name, version, target_model, created_at FROM prompts WHERE name LIKE ? ORDER BY created_at DESC",
            (f"%{name_filter}%",),
        )
    else:
        cur.execute("SELECT id, name, version, target_model, created_at FROM prompts ORDER BY created_at DESC")
    rows = cur.fetchall()
    conn.close()
    return [dict(r) for r in rows]


def delete_prompt(prompt_id: int) -> bool:
    conn = get_conn()
    cur = conn.cursor()
    cur.execute("DELETE FROM versions WHERE prompt_id = ?", (prompt_id,))
    cur.execute("DELETE FROM prompts WHERE id = ?", (prompt_id,))
    conn.commit()
    changed = cur.rowcount > 0
    conn.close()
    return changed


def get_versions(prompt_id: int) -> list:
    conn = get_conn()
    cur = conn.cursor()
    cur.execute("SELECT * FROM versions WHERE prompt_id = ? ORDER BY created_at ASC", (prompt_id,))
    rows = cur.fetchall()
    conn.close()
    return [dict(r) for r in rows]


def get_prompts_for_compare(id1: int, id2: int) -> tuple:
    p1 = get_prompt(id1)
    p2 = get_prompt(id2)
    return (p1, p2)