|
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} |
|
|