Spaces:
Configuration error
Configuration error
File size: 3,684 Bytes
ffd392f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 | """
db.py - SQLite storage layer for Local AI.
Uses Python's built-in sqlite3 module (zero extra dependencies, fully offline).
"""
import sqlite3
import json
import os
import time
import hashlib
from contextlib import contextmanager
from typing import Optional
DB_PATH = os.path.join(os.path.dirname(os.path.dirname(__file__)), "data", "local_ai.db")
SCHEMA = """
CREATE TABLE IF NOT EXISTS documents (
id INTEGER PRIMARY KEY AUTOINCREMENT,
filename TEXT NOT NULL,
file_hash TEXT,
doc_type TEXT,
raw_text TEXT,
extracted_json TEXT,
confidence REAL,
processing_time REAL,
created_at REAL,
pipeline_mode TEXT
);
CREATE INDEX IF NOT EXISTS idx_doc_type ON documents(doc_type);
CREATE INDEX IF NOT EXISTS idx_file_hash ON documents(file_hash);
"""
@contextmanager
def get_conn():
os.makedirs(os.path.dirname(DB_PATH), exist_ok=True)
conn = sqlite3.connect(DB_PATH)
conn.row_factory = sqlite3.Row
try:
yield conn
conn.commit()
finally:
conn.close()
def init_db():
with get_conn() as conn:
conn.executescript(SCHEMA)
def file_hash_bytes(data: bytes) -> str:
return hashlib.sha256(data).hexdigest()
def find_by_hash(file_hash: str) -> Optional[sqlite3.Row]:
with get_conn() as conn:
cur = conn.execute("SELECT * FROM documents WHERE file_hash = ? LIMIT 1", (file_hash,))
return cur.fetchone()
def insert_document(filename, file_hash, doc_type, raw_text, extracted: dict,
confidence: float, processing_time: float, pipeline_mode: str) -> int:
with get_conn() as conn:
cur = conn.execute(
"""INSERT INTO documents
(filename, file_hash, doc_type, raw_text, extracted_json,
confidence, processing_time, created_at, pipeline_mode)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)""",
(filename, file_hash, doc_type, raw_text, json.dumps(extracted),
confidence, processing_time, time.time(), pipeline_mode),
)
return cur.lastrowid
def list_documents(limit: int = 200):
with get_conn() as conn:
cur = conn.execute(
"SELECT * FROM documents ORDER BY created_at DESC LIMIT ?", (limit,)
)
return cur.fetchall()
def get_document(doc_id: int):
with get_conn() as conn:
cur = conn.execute("SELECT * FROM documents WHERE id = ?", (doc_id,))
return cur.fetchone()
def delete_document(doc_id: int):
with get_conn() as conn:
conn.execute("DELETE FROM documents WHERE id = ?", (doc_id,))
def search_documents(query: str):
q = f"%{query.lower()}%"
with get_conn() as conn:
cur = conn.execute(
"""SELECT * FROM documents
WHERE lower(filename) LIKE ?
OR lower(raw_text) LIKE ?
OR lower(extracted_json) LIKE ?
ORDER BY created_at DESC""",
(q, q, q),
)
return cur.fetchall()
def stats():
with get_conn() as conn:
total = conn.execute("SELECT COUNT(*) c FROM documents").fetchone()["c"]
avg_conf = conn.execute("SELECT AVG(confidence) c FROM documents").fetchone()["c"] or 0
avg_time = conn.execute("SELECT AVG(processing_time) c FROM documents").fetchone()["c"] or 0
by_type = conn.execute(
"SELECT doc_type, COUNT(*) c FROM documents GROUP BY doc_type"
).fetchall()
return {
"total_documents": total,
"avg_confidence": round(avg_conf, 3),
"avg_processing_time": round(avg_time, 3),
"by_type": {r["doc_type"]: r["c"] for r in by_type},
}
|