File size: 2,242 Bytes
5fffd14 6950cd1 5fffd14 6950cd1 5fffd14 6950cd1 5fffd14 |
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 |
import os
import json
from datetime import datetime
from typing import List, Dict, Any
class SimpleDB:
def __init__(self, db_path: str = "./data/documents.json"):
"""Initialize a simple JSON-based document store"""
self.db_path = db_path
os.makedirs(os.path.dirname(db_path), exist_ok=True)
# Create the DB file if it doesn't exist
if not os.path.exists(db_path):
with open(db_path, 'w') as f:
json.dump({"documents": [], "queries": []}, f)
def _read_db(self):
"""Read the database file"""
with open(self.db_path, 'r') as f:
return json.load(f)
def _write_db(self, data):
"""Write to the database file"""
with open(self.db_path, 'w') as f:
json.dump(data, f, indent=2)
def add_document(self, filename: str, file_path: str, file_type: str):
"""Add a document to the database"""
db = self._read_db()
# Generating a simple ID
doc_id = len(db["documents"]) + 1
# Add document
db["documents"].append({
"id": doc_id,
"filename": filename,
"file_path": file_path,
"file_type": file_type,
"upload_date": str(datetime.now())
})
self._write_db(db)
return doc_id
def get_document(self, doc_id: int):
"""Get document by ID"""
db = self._read_db()
for doc in db["documents"]:
if doc["id"] == doc_id:
return doc
return None
def get_all_documents(self):
"""Get all documents"""
db = self._read_db()
return db["documents"]
def log_query(self, query_text: str, response: str = None):
"""Log a user query and its response"""
db = self._read_db()
# Generating a simple ID
query_id = len(db["queries"]) + 1
# Adding query
db["queries"].append({
"id": query_id,
"query_text": query_text,
"response": response,
"timestamp": str(datetime.now())
})
self._write_db(db)
return query_id |