Spaces:
Sleeping
Sleeping
File size: 3,360 Bytes
b2b931e 6f1c390 b2b931e 6f1c390 b2b931e 8fda2b2 b2b931e f93e387 b2b931e 6f1c390 f93e387 6f1c390 b2b931e f93e387 b2b931e 6f1c390 b2b931e 6f1c390 b2b931e 6f1c390 | 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 | import sqlite3
import pandas as pd
from datetime import datetime
import os
DB_NAME = "navy_docs.db"
class DatabaseManager:
def __init__(self, db_path=DB_NAME):
self.db_path = db_path
self._init_db()
def _init_db(self):
"""Creates tables if they don't exist."""
conn = sqlite3.connect(self.db_path)
c = conn.cursor()
# UPDATED: Added 'abstract' column
c.execute('''
CREATE TABLE IF NOT EXISTS documents (
doc_id TEXT PRIMARY KEY,
filename TEXT,
upload_date TEXT,
user_id TEXT,
full_text TEXT,
abstract TEXT
)
''')
# (Chunks table remains the same...)
c.execute('''
CREATE TABLE IF NOT EXISTS chunks (
chunk_id TEXT PRIMARY KEY,
doc_id TEXT,
text TEXT,
chunk_index INTEGER,
FOREIGN KEY(doc_id) REFERENCES documents(doc_id)
)
''')
conn.commit()
conn.close()
# UPDATED: Added abstract argument
def add_document(self, doc_id, filename, full_text, abstract="", user_id="default"):
conn = sqlite3.connect(self.db_path)
c = conn.cursor()
c.execute(
"INSERT OR REPLACE INTO documents (doc_id, filename, upload_date, user_id, full_text, abstract) VALUES (?, ?, ?, ?, ?, ?)",
(doc_id, filename, datetime.now().isoformat(), user_id, full_text, abstract)
)
conn.commit()
conn.close()
# NEW: Getter for the abstract
def get_doc_abstract(self, doc_id):
conn = sqlite3.connect(self.db_path)
c = conn.cursor()
c.execute("SELECT abstract FROM documents WHERE doc_id=?", (doc_id,))
result = c.fetchone()
conn.close()
return result[0] if result else "No abstract available."
def get_doc_text(self, doc_id):
conn = sqlite3.connect(self.db_path)
c = conn.cursor()
c.execute("SELECT full_text FROM documents WHERE doc_id=?", (doc_id,))
result = c.fetchone()
conn.close()
return result[0] if result else "Error: Document text not found."
def get_all_filenames(self):
conn = sqlite3.connect(self.db_path)
c = conn.cursor()
c.execute("SELECT filename FROM documents")
results = [r[0] for r in c.fetchall()]
conn.close()
return results
def delete_document(self, filename):
"""Deletes a document and its chunks by filename."""
conn = sqlite3.connect(self.db_path)
c = conn.cursor()
# 1. Get doc_id
c.execute("SELECT doc_id FROM documents WHERE filename=?", (filename,))
res = c.fetchone()
if not res:
conn.close()
return None
doc_id = res[0]
# 2. Delete chunks first (Foreign Key hygiene)
c.execute("DELETE FROM chunks WHERE doc_id=?", (doc_id,))
# 3. Delete parent
c.execute("DELETE FROM documents WHERE doc_id=?", (doc_id,))
conn.commit()
conn.close()
return doc_id |