Spaces:
Running on Zero
Running on Zero
| from __future__ import annotations | |
| import shutil | |
| import uuid | |
| from pathlib import Path | |
| from typing import Any, Dict, List, Optional | |
| from langchain_core.documents import Document | |
| from langchain_core.chat_history import InMemoryChatMessageHistory | |
| from src.chunking import get_chunker | |
| from src.config import settings | |
| from src.embedding_model import get_embeddings | |
| from src.logger import logger | |
| from src.prompt import RAG_SYSTEM_MESSAGE | |
| from src.guardrails import run_input_guardrails, run_output_guardrails | |
| from src.indexing import build_index | |
| from src.llm_client import get_openai_response_with_trace | |
| class SessionState: | |
| def __init__(self, session_id: str): | |
| self.session_id = session_id | |
| self.mode: str = "none" # none | direct | rag | |
| self.filename: Optional[str] = None | |
| self.file_size: int = 0 | |
| self.direct_text: Optional[str] = None | |
| self.vectorstore = None | |
| self.collection_name: Optional[str] = None | |
| self.chunk_count: int = 0 | |
| self.persist_dir: Optional[Path] = None | |
| self.pages: int = 0 | |
| self.chunks: List[Document] = [] | |
| self.memory = InMemoryChatMessageHistory() | |
| self.llm_provider: str = "openai" | |
| self.llm_model: str = getattr(settings, "model_name", "gpt-4o-mini") | |
| self.api_key: Optional[str] = None | |
| self.api_base_url: Optional[str] = None | |
| self.llm_client = None | |
| self.embedding_model_name: Optional[str] = None | |
| self.vector_db_type_used: Optional[str] = None | |
| self.chunk_strategy_used: Optional[str] = None | |
| self.enable_guardrails: bool = True | |
| self.guardrail_injection: bool = True | |
| self.guardrail_pii: bool = True | |
| self.guardrail_numeric: bool = True | |
| class RAGEngine: | |
| def __init__(self): | |
| self.sessions: Dict[str, SessionState] = {} | |
| logger.info("Initializing RAG engine") | |
| try: | |
| self._embeddings = get_embeddings() | |
| logger.info( | |
| "Embeddings ready", | |
| context={"model": settings.embedding_model}, | |
| ) | |
| except Exception as e: | |
| logger.error("Failed to initialize embeddings", context={"error": str(e)}) | |
| self._embeddings = None | |
| # No global LLM client required at startup (API keys come per session) | |
| self._client = None | |
| # ------------------------------------------------------------------ | |
| # Session management | |
| # ------------------------------------------------------------------ | |
| def create_session(self) -> str: | |
| session_id = str(uuid.uuid4()) | |
| self.sessions[session_id] = SessionState(session_id) | |
| logger.info("Session created", context={"session_id": session_id}) | |
| return session_id | |
| # Replace the entire get_session method in RAGEngine class: | |
| def get_session(self, session_id: str) -> Optional[SessionState]: | |
| state = self.sessions.get(session_id) | |
| if state is not None: | |
| return state | |
| import pickle | |
| from langchain_community.vectorstores import FAISS | |
| # Try to recover from FAISS first (since that's what you're using) | |
| target = Path(settings.faiss_persist_directory) / session_id | |
| chunks_path = target / "chunks.pkl" | |
| meta_path = target / "session_meta.pkl" | |
| if chunks_path.exists(): | |
| state = SessionState(session_id) | |
| try: | |
| with open(chunks_path, "rb") as f: | |
| state.chunks = pickle.load(f) | |
| except Exception as e: | |
| logger.error("Failed to load chunks", context={"error": str(e)}) | |
| return None | |
| meta = {} | |
| if meta_path.exists(): | |
| try: | |
| with open(meta_path, "rb") as f: | |
| meta = pickle.load(f) | |
| except Exception as e: | |
| logger.error("Failed to load metadata", context={"error": str(e)}) | |
| state.mode = meta.get("mode", "rag") | |
| state.filename = meta.get("filename") | |
| state.vector_db_type_used = meta.get("vector_db_type", "faiss") | |
| state.embedding_model_name = meta.get("embedding_model") | |
| state.chunk_count = len(state.chunks) | |
| state.persist_dir = target | |
| # Restore vectorstore | |
| try: | |
| emb = get_embeddings(state.embedding_model_name or settings.embedding_model) | |
| if state.vector_db_type_used == "faiss": | |
| state.vectorstore = FAISS.load_local( | |
| str(target), | |
| emb, | |
| allow_dangerous_deserialization=True, | |
| ) | |
| except Exception as e: | |
| logger.error("Failed to reload vectorstore", context={"error": str(e)}) | |
| state.vectorstore = None | |
| self.sessions[session_id] = state | |
| logger.info("Recovered session from disk", context={"session_id": session_id}) | |
| return state | |
| # Also check for Chroma recovery (fallback) | |
| target_chroma = Path(settings.chroma_persist_directory) / session_id | |
| chunks_path_chroma = target_chroma / "chunks.pkl" | |
| meta_path_chroma = target_chroma / "session_meta.pkl" | |
| if chunks_path_chroma.exists(): | |
| state = SessionState(session_id) | |
| try: | |
| with open(chunks_path_chroma, "rb") as f: | |
| state.chunks = pickle.load(f) | |
| except Exception as e: | |
| logger.error("Failed to load chunks from chroma", context={"error": str(e)}) | |
| return None | |
| meta = {} | |
| if meta_path_chroma.exists(): | |
| try: | |
| with open(meta_path_chroma, "rb") as f: | |
| meta = pickle.load(f) | |
| except Exception as e: | |
| logger.error("Failed to load metadata from chroma", context={"error": str(e)}) | |
| state.mode = meta.get("mode", "rag") | |
| state.filename = meta.get("filename") | |
| state.vector_db_type_used = meta.get("vector_db_type", "chroma") | |
| state.embedding_model_name = meta.get("embedding_model") | |
| state.chunk_count = len(state.chunks) | |
| state.persist_dir = target_chroma | |
| self.sessions[session_id] = state | |
| logger.info("Recovered session from disk (chroma)", context={"session_id": session_id}) | |
| return state | |
| return None | |
| def delete_session(self, session_id: str) -> bool: | |
| state = self.sessions.pop(session_id, None) | |
| if state is None: | |
| return False | |
| if state.persist_dir and state.persist_dir.exists(): | |
| try: | |
| shutil.rmtree(state.persist_dir, ignore_errors=True) | |
| logger.info("Cleaned persist dir", context={"path": str(state.persist_dir)}) | |
| except Exception as e: | |
| logger.error("Failed to clean persist dir", context={"error": str(e)}) | |
| logger.info("Session deleted", context={"session_id": session_id}) | |
| return True | |
| def list_sessions(self) -> List[str]: | |
| return list(self.sessions.keys()) | |
| # ------------------------------------------------------------------ | |
| # Ingestion | |
| # ------------------------------------------------------------------ | |
| def ingest_pdf( | |
| self, | |
| session_id: str, | |
| file_path: Path, | |
| original_filename: str, | |
| chunk_strategy: Optional[str] = None, | |
| embedding_model: Optional[str] = None, | |
| vector_db_type: Optional[str] = None, | |
| llm_provider: Optional[str] = None, | |
| llm_model: Optional[str] = None, | |
| api_key: Optional[str] = None, | |
| api_base_url: Optional[str] = None, | |
| enable_guardrails: Optional[bool] = None, | |
| guardrail_injection: Optional[bool] = None, | |
| guardrail_pii: Optional[bool] = None, | |
| guardrail_numeric: Optional[bool] = None, | |
| ) -> Dict[str, Any]: | |
| import pickle | |
| from src.llm_client import get_client_for_provider | |
| state = self.get_session(session_id) | |
| if state is None: | |
| raise ValueError(f"Unknown session: {session_id}") | |
| if llm_provider: | |
| state.llm_provider = llm_provider.lower() | |
| if llm_model: | |
| state.llm_model = llm_model | |
| if api_key: | |
| state.api_key = api_key | |
| if api_base_url is not None: | |
| state.api_base_url = api_base_url.strip() or None | |
| if embedding_model: | |
| state.embedding_model_name = embedding_model | |
| if enable_guardrails is not None: | |
| state.enable_guardrails = enable_guardrails | |
| if guardrail_injection is not None: | |
| state.guardrail_injection = guardrail_injection | |
| if guardrail_pii is not None: | |
| state.guardrail_pii = guardrail_pii | |
| if guardrail_numeric is not None: | |
| state.guardrail_numeric = guardrail_numeric | |
| state.llm_client = get_client_for_provider( | |
| state.llm_provider, | |
| state.api_key, | |
| base_url=state.api_base_url, | |
| ) | |
| emb_name = embedding_model or state.embedding_model_name or settings.embedding_model | |
| emb = get_embeddings(emb_name) | |
| steps: List[str] = [] | |
| steps.append(f"LLM: {state.llm_provider} / {state.llm_model}") | |
| steps.append(f"Embedding model: {emb_name}") | |
| file_size = file_path.stat().st_size | |
| state.filename = original_filename | |
| state.file_size = file_size | |
| # ------------------------------------------------------------------ | |
| # 1. Load PDF | |
| # ------------------------------------------------------------------ | |
| if getattr(settings, "enable_multimodal", False): | |
| from src.pdf_processor import load_pdf_multimodal | |
| from src.llm_client import describe_image_with_vision | |
| pages: List[Document] = load_pdf_multimodal( | |
| file_path, | |
| vision_fn=describe_image_with_vision, | |
| ) | |
| steps.append("Loaded PDF with multimodal processor (text + tables + vision)") | |
| else: | |
| from langchain_community.document_loaders import PyPDFLoader | |
| pages = PyPDFLoader(str(file_path)).load() | |
| steps.append("Loaded PDF with PyPDFLoader") | |
| if not pages: | |
| raise ValueError("PDF is empty or unreadable") | |
| state.pages = len(pages) | |
| full_text = "\n\n".join(p.page_content for p in pages) | |
| steps.append(f"Loaded PDF ({len(pages)} pages)") | |
| # ------------------------------------------------------------------ | |
| # 2. Direct context for small files | |
| # ------------------------------------------------------------------ | |
| if file_size <= settings.direct_context_threshold_bytes: | |
| state.mode = "direct" | |
| state.direct_text = full_text | |
| state.chunk_count = 0 | |
| state.chunks = [] | |
| state.vectorstore = None | |
| state.collection_name = None | |
| state.persist_dir = None | |
| steps.append("Small file → direct context mode (no indexing)") | |
| # Persist minimal meta so ZeroGPU chat worker can recover | |
| meta_dir = Path(settings.upload_path) / session_id | |
| meta_dir.mkdir(parents=True, exist_ok=True) | |
| with open(meta_dir / "session_meta.pkl", "wb") as f: | |
| pickle.dump( | |
| { | |
| "mode": "direct", | |
| "filename": original_filename, | |
| "direct_text": full_text, | |
| "embedding_model": emb_name, | |
| }, | |
| f, | |
| ) | |
| return { | |
| "session_id": session_id, | |
| "mode": "direct", | |
| "filename": original_filename, | |
| "file_size_bytes": file_size, | |
| "pages": len(pages), | |
| "chunks": 0, | |
| "message": ( | |
| f"Small file processed in direct context mode " | |
| f"(≤ {settings.direct_context_threshold_kb} KB)." | |
| ), | |
| "pipeline_steps": steps, | |
| "chunk_strategy_used": None, | |
| "embedding_model_used": emb_name, | |
| "vector_db_type_used": None, | |
| } | |
| # ------------------------------------------------------------------ | |
| # 3. Chunk | |
| # ------------------------------------------------------------------ | |
| state.mode = "rag" | |
| strategy = (chunk_strategy or settings.chunk_strategy).lower() | |
| state.chunk_strategy_used = strategy | |
| steps.append(f"Chunk strategy: {strategy}") | |
| splitter = get_chunker(strategy=strategy) | |
| chunks = splitter.split_documents(pages) | |
| state.chunk_count = len(chunks) | |
| state.chunks = chunks | |
| steps.append(f"Created {len(chunks)} chunks") | |
| # ------------------------------------------------------------------ | |
| # 4. Index | |
| # ------------------------------------------------------------------ | |
| index_type = (vector_db_type or settings.vector_db_type).lower() | |
| state.vector_db_type_used = index_type | |
| steps.append(f"Vector DB: {index_type}") | |
| vector_store = build_index( | |
| documents=chunks, | |
| embedding_model=emb, | |
| session_id=session_id, | |
| filename=original_filename, | |
| index_type=index_type, | |
| ) | |
| state.vectorstore = vector_store | |
| steps.append(f"Indexed into {index_type}") | |
| if index_type == "chroma": | |
| state.persist_dir = Path(settings.chroma_persist_directory) / session_id | |
| state.collection_name = f"session_{session_id}" | |
| elif index_type == "faiss": | |
| state.persist_dir = Path(settings.faiss_persist_directory) / session_id | |
| state.collection_name = f"session_{session_id}" | |
| else: | |
| state.persist_dir = Path(settings.upload_path) / session_id | |
| state.collection_name = None | |
| # ------------------------------------------------------------------ | |
| # 5. Persist chunks + meta for ZeroGPU worker recovery ← STEP 2 | |
| # ------------------------------------------------------------------ | |
| meta_dir = state.persist_dir or (Path(settings.upload_path) / session_id) | |
| meta_dir.mkdir(parents=True, exist_ok=True) | |
| with open(meta_dir / "chunks.pkl", "wb") as f: | |
| pickle.dump(chunks, f) | |
| with open(meta_dir / "session_meta.pkl", "wb") as f: | |
| pickle.dump( | |
| { | |
| "mode": state.mode, | |
| "filename": original_filename, | |
| "vector_db_type": index_type, | |
| "embedding_model": emb_name, | |
| "chunk_strategy": strategy, | |
| "pages": state.pages, | |
| }, | |
| f, | |
| ) | |
| steps.append(f"Persisted session artifacts to {meta_dir}") | |
| return { | |
| "session_id": session_id, | |
| "mode": "rag", | |
| "filename": original_filename, | |
| "file_size_bytes": file_size, | |
| "pages": len(pages), | |
| "chunks": len(chunks), | |
| "message": "PDF ingested into vector store.", | |
| "pipeline_steps": steps, | |
| "chunk_strategy_used": strategy, | |
| "embedding_model_used": emb_name, | |
| "vector_db_type_used": index_type, | |
| } | |
| def process_uploaded_file( | |
| self, | |
| session_id: str, | |
| file_path: Path, | |
| original_filename: Optional[str] = None, | |
| chunk_strategy: Optional[str] = None, | |
| embedding_model: Optional[str] = None, | |
| vector_db_type: Optional[str] = None, | |
| llm_provider: Optional[str] = None, | |
| llm_model: Optional[str] = None, | |
| api_key: Optional[str] = None, | |
| api_base_url: Optional[str] = None, | |
| enable_guardrails: Optional[bool] = None, | |
| guardrail_injection: Optional[bool] = None, | |
| guardrail_pii: Optional[bool] = None, | |
| guardrail_numeric: Optional[bool] = None, | |
| ) -> Dict[str, Any]: | |
| return self.ingest_pdf( | |
| session_id, | |
| file_path, | |
| original_filename or file_path.name, | |
| chunk_strategy=chunk_strategy, | |
| embedding_model=embedding_model, | |
| vector_db_type=vector_db_type, | |
| llm_provider=llm_provider, | |
| llm_model=llm_model, | |
| api_key=api_key, | |
| api_base_url=api_base_url, | |
| enable_guardrails=enable_guardrails, | |
| guardrail_injection=guardrail_injection, | |
| guardrail_pii=guardrail_pii, | |
| guardrail_numeric=guardrail_numeric, | |
| ) | |
| # ------------------------------------------------------------------ | |
| # Query | |
| # ------------------------------------------------------------------ | |
| def query( | |
| self, | |
| session_id: str, | |
| question: str, | |
| k: Optional[int] = None, | |
| retrieval_type: Optional[str] = None, | |
| llm_provider: Optional[str] = None, | |
| llm_model: Optional[str] = None, | |
| api_key: Optional[str] = None, | |
| api_base_url: Optional[str] = None, | |
| ) -> Dict[str, Any]: | |
| from src.llm_client import get_client_for_provider | |
| state = self.get_session(session_id) | |
| if state is None: | |
| raise ValueError("Session not found") | |
| if state.mode == "none": | |
| raise ValueError("No document uploaded for this session") | |
| # ----- Per-request overrides (required on ZeroGPU: keys don't survive workers) ----- | |
| if llm_provider: | |
| state.llm_provider = llm_provider.lower() | |
| if llm_model: | |
| state.llm_model = llm_model | |
| if api_key: | |
| state.api_key = api_key | |
| if api_base_url is not None: | |
| state.api_base_url = (api_base_url or "").strip() or None | |
| steps: List[str] = [] | |
| client = get_client_for_provider( | |
| state.llm_provider, | |
| state.api_key, | |
| base_url=state.api_base_url, | |
| ) | |
| state.llm_client = client | |
| model = state.llm_model or settings.model_name | |
| rtype = ( | |
| retrieval_type | |
| or getattr(settings, "retrieval_type", None) | |
| or "semantic" | |
| ).lower() | |
| k = k or settings.num_retrieved_chunks | |
| steps.append(f"LLM: {state.llm_provider} / {model}") | |
| steps.append(f"Retrieval type: {rtype}") | |
| steps.append(f"Top-k: {k}") | |
| # ------------------------------------------------------------------ | |
| # 1. Input guardrails | |
| # ------------------------------------------------------------------ | |
| use_guardrails = getattr(state, "enable_guardrails", True) | |
| if use_guardrails: | |
| steps.append("Input guardrails: enabled") | |
| input_guard = run_input_guardrails(question) | |
| if not getattr(state, "guardrail_injection", True) and input_guard.injection_detected: | |
| input_guard.injection_detected = False | |
| input_guard.overall_passed = True | |
| input_guard.status = "passed" | |
| input_guard.details = [d for d in input_guard.details if "Injection" not in d] | |
| input_guard.details.append("Injection check: skipped by user") | |
| if not getattr(state, "guardrail_pii", True) and input_guard.pii_detected: | |
| input_guard.pii_detected = False | |
| input_guard.pii_entities = [] | |
| input_guard.redacted_query = None | |
| input_guard.warnings = [w for w in input_guard.warnings if "PII" not in w] | |
| input_guard.details.append("PII check: skipped by user") | |
| if input_guard.overall_passed and not input_guard.warnings: | |
| input_guard.status = "passed" | |
| if not input_guard.overall_passed: | |
| steps.append("Input guardrails: BLOCKED") | |
| return { | |
| "mode": "blocked", | |
| "answer": ( | |
| "Your question was blocked by safety guardrails.\n" | |
| + "\n".join(input_guard.details) | |
| ), | |
| "sources": [], | |
| "guardrails": { | |
| "input": input_guard.to_dict(), | |
| "output": None, | |
| "status": "blocked", | |
| "warnings": input_guard.warnings, | |
| "ungrounded_numbers": [], | |
| }, | |
| "reasoning_trace": None, | |
| "pipeline_steps": steps, | |
| "retrieval_type_used": rtype, | |
| "history_turns": len(state.memory.messages) // 2, | |
| } | |
| safe_question = input_guard.redacted_query or question | |
| if input_guard.warnings: | |
| steps.append(f"Input guardrails: warned ({len(input_guard.warnings)})") | |
| else: | |
| steps.append("Input guardrails: passed") | |
| else: | |
| steps.append("Input guardrails: skipped") | |
| class _Pass: | |
| overall_passed = True | |
| redacted_query = question | |
| warnings: List[str] = [] | |
| details: List[str] = [] | |
| ungrounded_numbers: List[str] = [] | |
| def to_dict(self): | |
| return { | |
| "overall_passed": True, | |
| "injection_detected": False, | |
| "pii_detected": False, | |
| "pii_entities": [], | |
| "warnings": [], | |
| "details": ["Guardrails disabled for this session"], | |
| } | |
| input_guard = _Pass() | |
| safe_question = question | |
| # ------------------------------------------------------------------ | |
| # 2. Retrieve context | |
| # ------------------------------------------------------------------ | |
| from src.retrieval import get_retriever, retrieve_documents | |
| if state.mode == "direct": | |
| context = state.direct_text or "" | |
| context_texts = [context] if context else [] | |
| sources: List[Dict[str, Any]] = [] | |
| steps.append("Mode: direct context (no retrieval)") | |
| else: | |
| if state.vectorstore is None and rtype in ("semantic", "hybrid"): | |
| if rtype != "bm25": | |
| raise ValueError( | |
| "Vector store not available for semantic/hybrid retrieval" | |
| ) | |
| if not state.chunks: | |
| raise ValueError("Chunks not available") | |
| retriever = get_retriever( | |
| vectorstore=state.vectorstore, | |
| documents=state.chunks, | |
| k=k, | |
| retrieval_type=rtype, | |
| ) | |
| docs = retrieve_documents(retriever, safe_question, k=k) | |
| context_texts = [doc.page_content for doc in docs] | |
| context = "\n\n".join( | |
| f"[Chunk {i + 1}]\n{doc.page_content}" for i, doc in enumerate(docs) | |
| ) | |
| sources = [ | |
| { | |
| "content": doc.page_content[:300] | |
| + ("..." if len(doc.page_content) > 300 else ""), | |
| "metadata": doc.metadata, | |
| } | |
| for doc in docs | |
| ] | |
| steps.append(f"Retrieved {len(docs)} chunks") | |
| # ------------------------------------------------------------------ | |
| # 3. Prompt | |
| # ------------------------------------------------------------------ | |
| history_text = "" | |
| if state.memory.messages: | |
| history_text = "Previous conversation:\n" | |
| for msg in state.memory.messages: | |
| role = "User" if msg.type == "human" else "Assistant" | |
| history_text += f"{role}: {msg.content}\n" | |
| history_text += "\n" | |
| steps.append(f"History turns: {len(state.memory.messages) // 2}") | |
| full_input = ( | |
| f"{history_text}" | |
| f"<Context>\n{context}\n</Context>\n" | |
| f"<Question>{safe_question}</Question>" | |
| ) | |
| # ------------------------------------------------------------------ | |
| # 4. LLM | |
| # ------------------------------------------------------------------ | |
| steps.append("Calling LLM…") | |
| answer, reasoning_trace = get_openai_response_with_trace( | |
| input=full_input, | |
| instructions=RAG_SYSTEM_MESSAGE, | |
| client=client, | |
| model=model, | |
| provider=state.llm_provider, | |
| reasoning_effort=getattr(settings, "reasoning_effort", None), | |
| ) | |
| steps.append("LLM response received") | |
| # ------------------------------------------------------------------ | |
| # 5. Output guardrails | |
| # ------------------------------------------------------------------ | |
| if use_guardrails and getattr(state, "guardrail_numeric", True): | |
| steps.append("Output guardrails: numeric grounding") | |
| output_guard = run_output_guardrails(answer, context_texts) | |
| if output_guard.warnings: | |
| steps.append(f"Output guardrails: warned ({len(output_guard.warnings)})") | |
| else: | |
| steps.append("Output guardrails: passed") | |
| else: | |
| steps.append("Output guardrails: skipped") | |
| class _OutPass: | |
| warnings: List[str] = [] | |
| ungrounded_numbers: List[str] = [] | |
| details: List[str] = [] | |
| def to_dict(self): | |
| return { | |
| "numeric_grounding_passed": True, | |
| "ungrounded_numbers": [], | |
| "warnings": [], | |
| "details": ["Output guardrails disabled"], | |
| } | |
| output_guard = _OutPass() | |
| # ------------------------------------------------------------------ | |
| # 6. Memory | |
| # ------------------------------------------------------------------ | |
| state.memory.add_user_message(question) | |
| state.memory.add_ai_message(answer) | |
| max_messages = settings.max_history_turns * 2 | |
| if len(state.memory.messages) > max_messages: | |
| state.memory.messages = state.memory.messages[-max_messages:] | |
| combined = { | |
| "input": input_guard.to_dict(), | |
| "output": output_guard.to_dict(), | |
| "status": ( | |
| "blocked" | |
| if not getattr(input_guard, "overall_passed", True) | |
| else ( | |
| "warned" | |
| if (input_guard.warnings or output_guard.warnings) | |
| else "passed" | |
| ) | |
| ), | |
| "warnings": list(input_guard.warnings) + list(output_guard.warnings), | |
| "ungrounded_numbers": list( | |
| getattr(output_guard, "ungrounded_numbers", []) or [] | |
| ), | |
| } | |
| return { | |
| "mode": state.mode, | |
| "answer": answer, | |
| "sources": sources, | |
| "guardrails": combined, | |
| "reasoning_trace": reasoning_trace or None, | |
| "pipeline_steps": steps, | |
| "retrieval_type_used": rtype, | |
| "history_turns": len(state.memory.messages) // 2, | |
| } | |
| def get_context(self, session_id: str) -> str: | |
| state = self.get_session(session_id) | |
| if state is None: | |
| raise ValueError("Session not found") | |
| if state.mode == "direct": | |
| return state.direct_text or "" | |
| return "" | |
| def get_status(self, session_id: str) -> Dict[str, Any]: | |
| state = self.get_session(session_id) | |
| if state is None: | |
| raise ValueError("Session not found") | |
| return { | |
| "session_id": session_id, | |
| "mode": state.mode, | |
| "filename": state.filename, | |
| "file_size": state.file_size, | |
| "chunk_count": state.chunk_count, | |
| "pages": state.pages, | |
| } | |
| def clear_history(self, session_id: str) -> bool: | |
| state = self.get_session(session_id) | |
| if state is None: | |
| return False | |
| state.memory.clear() | |
| logger.info("History cleared", context={"session_id": session_id}) | |
| return True | |
| rag_engine = RAGEngine() |