File size: 1,538 Bytes
5306da4 |
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 |
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}
|