llms-py 3.0.0__py3-none-any.whl → 3.0.0b1__py3-none-any.whl

llms/extensions/app/db.py

@@ -1,644 +0,0 @@
-import json
-import os
-import threading
-from datetime import datetime, timedelta
-from typing import Any, Dict
-
-from .db_manager import DbManager
-
-
-def with_user(data, user):
-    if user is None:
-        if "user" in data:
-            del data["user"]
-        return data
-    else:
-        data["user"] = user
-        return data
-
-
-def valid_columns(all_columns, fields):
-    if fields:
-        if not isinstance(fields, list):
-            fields = fields.split(",")
-        cols = []
-        for k in fields:
-            k = k.strip()
-            if k in all_columns:
-                cols.append(k)
-        return cols
-    return []
-
-
-def table_columns(all_columns, fields):
-    cols = valid_columns(all_columns, fields)
-    return ", ".join(cols) if len(cols) > 0 else ", ".join(all_columns)
-
-
-def select_columns(all_columns, fields, select=None):
-    columns = table_columns(all_columns, fields)
-    if select == "distinct":
-        return f"SELECT DISTINCT {columns}"
-    return f"SELECT {columns}"
-
-
-def order_by(all_columns, sort):
-    cols = []
-    for k in sort.split(","):
-        k = k.strip()
-        by = ""
-        if k[0] == "-":
-            by = " DESC"
-            k = k[1:]
-        if k in all_columns:
-            cols.append(f"{k}{by}")
-    return f"ORDER BY {', '.join(cols)} " if len(cols) > 0 else ""
-
-
-class AppDB:
-    def __init__(self, ctx, db_path):
-        if db_path is None:
-            raise ValueError("db_path is required")
-
-        self.ctx = ctx
-        self.db_path = str(db_path)
-
-        dirname = os.path.dirname(self.db_path)
-        if dirname:
-            os.makedirs(dirname, exist_ok=True)
-
-        self.db = DbManager(ctx, self.db_path)
-        self.columns = {
-            "thread": {
-                "id": "INTEGER",
-                "user": "TEXT",
-                "createdAt": "TIMESTAMP",
-                "updatedAt": "TIMESTAMP",
-                "title": "TEXT",
-                "systemPrompt": "TEXT",
-                "model": "TEXT",
-                "modelInfo": "JSON",
-                "modalities": "JSON",
-                "messages": "JSON",
-                "args": "JSON",
-                "toolHistory": "JSON",
-                "cost": "REAL",
-                "inputTokens": "INTEGER",
-                "outputTokens": "INTEGER",
-                "stats": "JSON",
-                "provider": "TEXT",
-                "providerModel": "TEXT",
-                "publishedAt": "TIMESTAMP",
-                "startedAt": "TIMESTAMP",
-                "completedAt": "TIMESTAMP",
-                "metadata": "JSON",
-                "error": "TEXT",
-                "ref": "TEXT",
-            },
-            "request": {
-                "id": "INTEGER",
-                "user": "TEXT",
-                "threadId": "INTEGER",
-                "createdAt": "TIMESTAMP",
-                "updatedAt": "TIMESTAMP",
-                "title": "TEXT",
-                "model": "TEXT",
-                "duration": "INTEGER",
-                "cost": "REAL",
-                "inputPrice": "REAL",
-                "inputTokens": "INTEGER",
-                "inputCachedTokens": "INTEGER",
-                "outputPrice": "REAL",
-                "outputTokens": "INTEGER",
-                "totalTokens": "INTEGER",
-                "usage": "JSON",
-                "provider": "TEXT",
-                "providerModel": "TEXT",
-                "providerRef": "TEXT",
-                "finishReason": "TEXT",
-                "startedAt": "TIMESTAMP",
-                "completedAt": "TIMESTAMP",
-                "error": "TEXT",
-                "stackTrace": "TEXT",
-                "ref": "TEXT",
-            },
-        }
-        with self.create_writer_connection() as conn:
-            self.init_db(conn)
-
-    def get_connection(self):
-        return self.create_reader_connection()
-
-    def create_reader_connection(self):
-        return self.db.create_reader_connection()
-
-    def create_writer_connection(self):
-        return self.db.create_writer_connection()
-
-    # Check for missing columns and migrate if necessary
-    def add_missing_columns(self, conn, table):
-        cur = self.db.exec(conn, f"PRAGMA table_info({table})")
-        columns = {row[1] for row in cur.fetchall()}
-
-        for col, dtype in self.columns[table].items():
-            if col not in columns:
-                try:
-                    self.db.exec(conn, f"ALTER TABLE {table} ADD COLUMN {col} {dtype}")
-                except Exception as e:
-                    self.ctx.err(f"adding {table} column {col}", e)
-
-    def init_db(self, conn):
-        # Create table with all columns
-        # Note: default SQLite timestamp has different tz to datetime.now()
-        overrides = {
-            "id": "INTEGER PRIMARY KEY AUTOINCREMENT",
-            "createdAt": "TIMESTAMP DEFAULT CURRENT_TIMESTAMP",
-            "updatedAt": "TIMESTAMP DEFAULT CURRENT_TIMESTAMP",
-        }
-        sql_columns = ",".join([f"{col} {overrides.get(col, dtype)}" for col, dtype in self.columns["thread"].items()])
-        self.db.exec(
-            conn,
-            f"""
-            CREATE TABLE IF NOT EXISTS thread (
-                {sql_columns}
-            )
-            """,
-        )
-        self.add_missing_columns(conn, "thread")
-        self.db.exec(conn, "CREATE INDEX IF NOT EXISTS idx_thread_user ON thread(user)")
-        self.db.exec(conn, "CREATE INDEX IF NOT EXISTS idx_thread_createdat ON thread(createdAt)")
-        self.db.exec(conn, "CREATE INDEX IF NOT EXISTS idx_thread_updatedat ON thread(updatedAt)")
-        self.db.exec(conn, "CREATE INDEX IF NOT EXISTS idx_thread_model ON thread(model)")
-        self.db.exec(conn, "CREATE INDEX IF NOT EXISTS idx_thread_cost ON thread(cost)")
-
-        sql_columns = ",".join([f"{col} {overrides.get(col, dtype)}" for col, dtype in self.columns["request"].items()])
-        self.db.exec(
-            conn,
-            f"""
-            CREATE TABLE IF NOT EXISTS request (
-                {sql_columns}
-            )
-            """,
-        )
-        self.add_missing_columns(conn, "request")
-        self.db.exec(conn, "CREATE INDEX IF NOT EXISTS idx_request_user ON request(user)")
-        self.db.exec(conn, "CREATE INDEX IF NOT EXISTS idx_request_createdat ON request(createdAt)")
-        self.db.exec(conn, "CREATE INDEX IF NOT EXISTS idx_request_cost ON request(cost)")
-        self.db.exec(conn, "CREATE INDEX IF NOT EXISTS idx_request_threadid ON request(threadId)")
-
-    def import_db(self, threads, requests):
-        self.ctx.log("import threads and requests")
-        with self.create_writer_connection() as conn:
-            conn.execute("DROP TABLE IF EXISTS thread")
-            conn.execute("DROP TABLE IF EXISTS request")
-            self.init_db(conn)
-            thread_id_map = {}
-            for thread in threads:
-                thread_id = self.import_thread(conn, thread)
-                thread_id_map[thread["id"]] = thread_id
-            self.ctx.log(f"imported {len(threads)} threads")
-            for request in requests:
-                self.import_request(conn, request, thread_id_map)
-            self.ctx.log(f"imported {len(requests)} requests")
-
-    def import_date(self, date):
-        # "1765794035" or "2025-12-31T05:41:46.686Z" or "2026-01-02 05:00:16"
-        str = date or datetime.now().isoformat()
-        if isinstance(str, int):
-            return datetime.fromtimestamp(str)
-        if isinstance(str, float):
-            return datetime.fromtimestamp(str)
-        return (
-            datetime.strptime(str, "%Y-%m-%dT%H:%M:%S.%fZ")
-            if "T" in str
-            else datetime.strptime(str, "%Y-%m-%d %H:%M:%S")
-        )
-
-    def import_thread(self, conn, orig):
-        thread = orig.copy()
-        thread["refId"] = thread["id"]
-        del thread["id"]
-
-        info = thread.get("modelInfo", thread.get("info", {}))
-        created_at = self.import_date(thread.get("createdAt"))
-        thread["createdAt"] = created_at
-        if "updateAt" not in thread:
-            thread["updateAt"] = created_at
-        thread["modelInfo"] = info
-        if "modalities" not in thread:
-            if "modalities" in info:
-                modalities = info["modalities"]
-                if isinstance(modalities, dict):
-                    input = modalities.get("input", ["text"])
-                    output = modalities.get("output", ["text"])
-                    thread["modalities"] = list(set(input + output))
-                else:
-                    thread["modalities"] = modalities
-            else:
-                thread["modalities"] = ["text"]
-        if "provider" not in thread and "provider" in info:
-            thread["provider"] = info["provider"]
-        if "providerModel" not in thread and "id" in info:
-            thread["providerModel"] = info["id"]
-
-        stats = thread.get("stats", {})
-        if "inputTokens" not in thread and "inputTokens" in stats:
-            thread["inputTokens"] = stats["inputTokens"]
-        if "outputTokens" not in thread and "outputTokens" in stats:
-            thread["outputTokens"] = stats["outputTokens"]
-        if "cost" not in thread and "cost" in stats:
-            thread["cost"] = stats["cost"]
-        if "completedAt" not in thread:
-            thread["completedAt"] = created_at + timedelta(milliseconds=stats.get("duration", 0))
-
-        sql_columns = []
-        sql_params = []
-        columns = self.columns["thread"]
-        for col in columns:
-            if col == "id":
-                continue
-            sql_columns.append(col)
-            val = thread.get(col, None)
-            if columns[col] == "JSON" and val is not None:
-                val = json.dumps(val)
-            sql_params.append(val)
-
-        return conn.execute(
-            f"INSERT INTO thread ({', '.join(sql_columns)}) VALUES ({', '.join(['?'] * len(sql_params))})",
-            sql_params,
-        ).lastrowid
-
-    # run on startup
-    def import_request(self, conn, orig, id_map):
-        request = orig.copy()
-        del request["id"]
-        thread_id = request.get("threadId")
-        if thread_id:
-            request["threadId"] = id_map.get(thread_id, None)
-
-        created_at = self.import_date(request.get("created"))
-        request["createdAt"] = created_at
-        if "updateAt" not in request:
-            request["updateAt"] = created_at
-        if "completedAt" not in request:
-            request["completedAt"] = created_at + timedelta(milliseconds=request.get("duration", 0))
-
-        sql_columns = []
-        sql_params = []
-        columns = self.columns["request"]
-        for col in columns:
-            if col == "id":
-                continue
-            sql_columns.append(col)
-            val = request.get(col, None)
-            if columns[col] == "JSON" and val is not None:
-                val = json.dumps(val)
-            sql_params.append(val)
-
-        return conn.execute(
-            f"INSERT INTO request ({', '.join(sql_columns)}) VALUES ({', '.join(['?'] * len(sql_params))})",
-            sql_params,
-        ).lastrowid
-
-    def get_user_filter(self, user=None, params=None):
-        if user is None:
-            return "WHERE user IS NULL", params or {}
-        else:
-            args = params.copy() if params else {}
-            args.update({"user": user})
-            return "WHERE user = :user", args
-
-    def get_thread(self, id, user=None):
-        try:
-            sql_where, params = self.get_user_filter(user, {"id": id})
-            return self.db.one(f"SELECT * FROM thread {sql_where} AND id = :id", params)
-        except Exception as e:
-            self.ctx.err(f"get_thread ({id}, {user})", e)
-            return None
-
-    def get_thread_column(self, id, column, user=None):
-        if column not in self.columns["thread"]:
-            self.ctx.err(f"get_thread_column invalid column ({id}, {column}, {user})", None)
-            return None
-
-        try:
-            sql_where, params = self.get_user_filter(user, {"id": id})
-            return self.db.scalar(f"SELECT {column} FROM thread {sql_where} AND id = :id", params)
-        except Exception as e:
-            self.ctx.err(f"get_thread_column ({id}, {column}, {user})", e)
-            return None
-
-    def query_threads(self, query: Dict[str, Any], user=None):
-        try:
-            columns = self.columns["thread"]
-            all_columns = columns.keys()
-
-            take = min(int(query.get("take", "50")), 1000)
-            skip = int(query.get("skip", "0"))
-            sort = query.get("sort", "-id")
-
-            # always filter by user
-            sql_where, params = self.get_user_filter(user, {"take": take, "skip": skip})
-
-            filter = {}
-            for k in query:
-                if k in all_columns:
-                    filter[k] = query[k]
-                    params[k] = query[k]
-
-            if len(filter) > 0:
-                sql_where += " AND " + " AND ".join([f"{k} = :{k}" for k in filter])
-
-            if "null" in query:
-                cols = valid_columns(all_columns, query["null"])
-                if len(cols) > 0:
-                    sql_where += " AND " + " AND ".join([f"{k} IS NULL" for k in cols])
-
-            if "not_null" in query:
-                cols = valid_columns(all_columns, query.get("not_null"))
-                if len(cols) > 0:
-                    sql_where += " AND " + " AND ".join([f"{k} IS NOT NULL" for k in cols])
-
-            if "q" in query:
-                sql_where += " AND " if sql_where else "WHERE "
-                sql_where += "(title LIKE :q OR messages LIKE :q)"
-                params["q"] = f"%{query['q']}%"
-
-            sql = f"{select_columns(all_columns, query.get('fields'), select=query.get('select'))} FROM thread {sql_where} {order_by(all_columns, sort)} LIMIT :take OFFSET :skip"
-
-            if query.get("as") == "column":
-                return self.db.column(sql, params)
-            else:
-                return self.db.all(sql, params)
-
-        except Exception as e:
-            self.ctx.err(f"query_threads ({take}, {skip})", e)
-            return []
-
-    def insert(self, table, info, callback=None):
-        if not info:
-            raise Exception("info is required")
-
-        columns = self.columns[table]
-        args = {}
-        known_columns = columns.keys()
-        for k, val in info.items():
-            if k in known_columns and k != "id":
-                args[k] = self.db.value(val)
-
-        insert_keys = list(args.keys())
-        insert_body = ", ".join(insert_keys)
-        insert_values = ", ".join(["?" for _ in insert_keys])
-
-        sql = f"INSERT INTO {table} ({insert_body}) VALUES ({insert_values})"
-
-        self.db.write(sql, tuple(args[k] for k in insert_keys), callback)
-
-    async def insert_async(self, table, info):
-        event = threading.Event()
-
-        ret = [None]
-
-        def cb(lastrowid, rowcount, error=None):
-            nonlocal ret
-            if error:
-                raise error
-            ret[0] = lastrowid
-            event.set()
-
-        self.insert(table, info, cb)
-        event.wait()
-        return ret[0]
-
-    def update(self, table, info, callback=None):
-        if not info:
-            raise Exception("info is required")
-
-        columns = self.columns[table]
-        args = {}
-        known_columns = columns.keys()
-        for k, val in info.items():
-            if k in known_columns and k != "id":
-                args[k] = self.db.value(val)
-
-        update_keys = list(args.keys())
-        update_body = ", ".join([f"{k} = :{k}" for k in update_keys])
-
-        args["id"] = info["id"]
-        sql = f"UPDATE {table} SET {update_body} WHERE id = :id"
-
-        self.db.write(sql, args, callback)
-
-    async def update_async(self, table, info):
-        event = threading.Event()
-
-        ret = [None]
-
-        def cb(lastrowid, rowcount, error=None):
-            nonlocal ret
-            if error:
-                raise error
-            ret[0] = rowcount
-            event.set()
-
-        self.update(table, info, cb)
-        event.wait()
-        return ret[0]
-
-    def prepare_thread(self, thread, id=None, user=None):
-        now = datetime.now()
-        if id:
-            thread["id"] = id
-        else:
-            thread["createdAt"] = now
-        thread["updatedAt"] = now
-        if "messages" in thread:
-            for m in thread["messages"]:
-                self.ctx.cache_message_inline_data(m)
-        return with_user(thread, user=user)
-
-    def create_thread(self, thread: Dict[str, Any], user=None):
-        return self.insert("thread", self.prepare_thread(thread, user=user))
-
-    async def create_thread_async(self, thread: Dict[str, Any], user=None):
-        return await self.insert_async("thread", self.prepare_thread(thread, user=user))
-
-    def update_thread(self, id, thread: Dict[str, Any], user=None):
-        return self.update("thread", self.prepare_thread(thread, id, user=user))
-
-    async def update_thread_async(self, id, thread: Dict[str, Any], user=None):
-        return await self.update_async("thread", self.prepare_thread(thread, id, user=user))
-
-    def delete_thread(self, id, user=None, callback=None):
-        sql_where, params = self.get_user_filter(user, {"id": id})
-        self.db.write(f"DELETE FROM thread {sql_where} AND id = :id", params, callback)
-
-    def query_requests(self, query: Dict[str, Any], user=None):
-        try:
-            columns = self.columns["request"]
-            all_columns = columns.keys()
-
-            take = min(int(query.get("take", "50")), 10000)
-            skip = int(query.get("skip", 0))
-            sort = query.get("sort", "-id")
-
-            # always filter by user
-            sql_where, params = self.get_user_filter(user, {"take": take, "skip": skip})
-
-            filter = {}
-            for k in query:
-                if k in all_columns:
-                    filter[k] = query[k]
-                    params[k] = query[k]
-
-            if len(filter) > 0:
-                sql_where += " AND " + " AND ".join([f"{k} = :{k}" for k in filter])
-
-            if "null" in query:
-                cols = valid_columns(all_columns, query["null"])
-                if len(cols) > 0:
-                    sql_where += " AND " + " AND ".join([f"{k} IS NULL" for k in cols])
-
-            if "not_null" in query:
-                cols = valid_columns(all_columns, query.get("not_null"))
-                if len(cols) > 0:
-                    sql_where += " AND " + " AND ".join([f"{k} IS NOT NULL" for k in cols])
-
-            if "q" in query:
-                sql_where += " AND " if sql_where else "WHERE "
-                sql_where += "(title LIKE :q)"
-                params["q"] = f"%{query['q']}%"
-
-            if "month" in query:
-                sql_where += " AND strftime('%Y-%m', createdAt) = :month"
-                params["month"] = query["month"]
-
-            sql = f"{select_columns(all_columns, query.get('fields'), select=query.get('select'))} FROM request {sql_where} {order_by(all_columns, sort)}LIMIT :take OFFSET :skip"
-
-            if query.get("as") == "column":
-                return self.db.column(sql, params)
-            else:
-                return self.db.all(sql, params)
-        except Exception as e:
-            self.ctx.err(f"query_requests ({take}, {skip})", e)
-            return []
-
-    def get_request_summary(self, user=None):
-        try:
-            sql_where, params = self.get_user_filter(user)
-            # Use strftime to format date as YYYY-MM-DD
-            sql = f"""
-                SELECT
-                    strftime('%Y-%m-%d', createdAt) as date,
-                    count(id) as requests,
-                    sum(cost) as cost,
-                    sum(inputTokens) as inputTokens,
-                    sum(outputTokens) as outputTokens
-                FROM request
-                {sql_where}
-                GROUP BY date
-                ORDER BY date
-            """
-            return self.db.all(sql, params)
-        except Exception as e:
-            self.ctx.err(f"get_request_summary ({user})", e)
-            return []
-
-    def get_daily_request_summary(self, day, user=None):
-        try:
-            sql_where, params = self.get_user_filter(user)
-            # Add date filter
-            sql_where += " AND strftime('%Y-%m-%d', createdAt) = :day"
-            params["day"] = day
-
-            # Model aggregation
-            sql_model = f"""
-                SELECT
-                    model,
-                    count(id) as count,
-                    sum(cost) as cost,
-                    sum(duration) as duration,
-                    sum(inputTokens + outputTokens) as tokens,
-                    sum(inputTokens) as inputTokens,
-                    sum(outputTokens) as outputTokens
-                FROM request
-                {sql_where}
-                GROUP BY model
-            """
-            model_data = {}
-            for row in self.db.all(sql_model, params):
-                model_data[row["model"]] = {
-                    "cost": row["cost"] or 0,
-                    "count": row["count"],
-                    "duration": row["duration"] or 0,
-                    "tokens": row["tokens"] or 0,
-                    "inputTokens": row["inputTokens"] or 0,
-                    "outputTokens": row["outputTokens"] or 0,
-                }
-
-            # Provider aggregation
-            sql_provider = f"""
-                SELECT
-                    provider,
-                    count(id) as count,
-                    sum(cost) as cost,
-                    sum(duration) as duration,
-                    sum(inputTokens + outputTokens) as tokens,
-                    sum(inputTokens) as inputTokens,
-                    sum(outputTokens) as outputTokens
-                FROM request
-                {sql_where}
-                AND provider IS NOT NULL
-                GROUP BY provider
-            """
-            provider_data = {}
-            for row in self.db.all(sql_provider, params):
-                provider_data[row["provider"]] = {
-                    "cost": row["cost"] or 0,
-                    "count": row["count"],
-                    "duration": row["duration"] or 0,
-                    "tokens": row["tokens"] or 0,
-                    "inputTokens": row["inputTokens"] or 0,
-                    "outputTokens": row["outputTokens"] or 0,
-                }
-
-            return {"modelData": model_data, "providerData": provider_data}
-        except Exception as e:
-            self.ctx.err(f"get_daily_request_summary ({day}, {user})", e)
-            return {"modelData": {}, "providerData": {}}
-
-    def create_request(self, request: Dict[str, Any], user=None):
-        request["createdAt"] = request["updatedAt"] = datetime.now()
-        return self.insert("request", with_user(request, user=user))
-
-    async def create_request_async(self, request: Dict[str, Any], user=None):
-        request["createdAt"] = request["updatedAt"] = datetime.now()
-        return await self.insert_async("request", with_user(request, user=user))
-
-    def update_request(self, id, request: Dict[str, Any], user=None):
-        request["id"] = id
-        request["updatedAt"] = datetime.now()
-        return self.update("request", with_user(request, user=user))
-
-    async def update_request_async(self, id, request: Dict[str, Any], user=None):
-        request["id"] = id
-        request["updatedAt"] = datetime.now()
-        return await self.update_async("request", with_user(request, user=user))
-
-    def delete_request(self, id, user=None, callback=None):
-        sql_where, params = self.get_user_filter(user, {"id": id})
-        self.db.write(f"DELETE FROM request {sql_where} AND id = :id", params, callback)
-
-    def close(self):
-        self.db.close()
-
-        # complete all in progress tasks
-        with self.db.create_writer_connection() as conn:
-            conn.execute(
-                "UPDATE thread SET completedAt = :completedAt, error = :error WHERE completedAt IS NULL",
-                {"completedAt": datetime.now(), "error": "Server Shutdown"},
-            )
-            conn.execute(
-                "UPDATE request SET completedAt = :completedAt, error = :error WHERE completedAt IS NULL",
-                {"completedAt": datetime.now(), "error": "Server Shutdown"},
-            )