VocRT / rag_database.py
Anurag
version-2 initial version
5306da4
import pickle
import os
from fastapi import FastAPI
from pydantic import BaseModel
from typing import Dict, Any
app = FastAPI()
FILE_PATH = "rag_database.pkl"
if not os.path.exists(FILE_PATH):
with open(FILE_PATH, "wb") as file:
pickle.dump({}, file)
def save_document_to_entry(session_id, filename, file_path):
try:
try:
with open(file_path, "rb") as file:
stored_data = pickle.load(file)
except (FileNotFoundError, EOFError):
stored_data = {}
if session_id not in stored_data:
stored_data[session_id] = {
"files": []
}
stored_data[session_id]["files"].append(filename)
with open(file_path, "wb") as file:
pickle.dump(stored_data, file)
except Exception as e:
print(f"Error saving file data: {e}")
def get_document_from_file():
try:
with open(FILE_PATH, "rb") as file:
data = pickle.load(file)
return data.get(session_id, [])
except (FileNotFoundError, pickle.UnpicklingError) as e:
print(f"Error reading or parsing file: {e}")
return []
except Exception as e:
print(f"Unexpected error: {e}")
return []
async def save_document():
save_document_to_file(doc.session_id, doc.data, doc.filename)
return {"message": "Document saved successfully!"}
async def get_document():
documents = get_document_from_file(sessionId)
return {"session_id": sessionId, "documents": documents}