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