Spaces:
Sleeping
Sleeping
File size: 6,469 Bytes
b106f38 c86ca34 b106f38 c86ca34 b106f38 a58c437 b106f38 | 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 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 | """Application Shared State and Service Dependency Provider.
Manages singletons for the vector store, embedding service, and RAG pipeline.
"""
from functools import lru_cache
import logging
from pathlib import Path
from typing import Optional
from app.config import Settings, get_settings
from app.embeddings.embedding_service import EmbeddingService, get_embedding_service
from app.generation.llm import LLMServiceFactory
from app.ingestion.chunker import RecursiveChunker
from app.ingestion.pdf_loader import PDFLoader
from app.rag.pipeline import RAGPipeline
from app.retrieval.query_rewriter import QueryRewriter
from app.retrieval.reranker import CrossEncoderReranker
from app.retrieval.retriever import DocumentRetriever
from app.vectorstore.base import BaseVectorStore
from app.vectorstore.chroma_store import ChromaVectorStore
from app.vectorstore.faiss_store import FAISSVectorStore
logger = logging.getLogger(__name__)
class AppState:
"""Encapsulates active services and indices across the application lifetime."""
def __init__(self, settings: Optional[Settings] = None):
self.settings = settings or get_settings()
self.pdf_loader = PDFLoader()
self.chunker = RecursiveChunker(
chunk_size=self.settings.chunk_size,
chunk_overlap=self.settings.chunk_overlap,
)
# 1. Embedding Service
self.embedding_service = EmbeddingService(
model_name=self.settings.embedding_model
)
# 2. Vector Database
self.vector_store: BaseVectorStore = self._init_vector_store()
# 3. Retriever
self.retriever = DocumentRetriever(
vector_store=self.vector_store,
embedding_service=self.embedding_service,
similarity_threshold=self.settings.similarity_threshold,
default_top_k=self.settings.top_k,
)
# 4. LLM Service
self.llm_service = LLMServiceFactory.create(
provider=self.settings.llm_provider,
gemini_api_key=self.settings.gemini_api_key,
nvidia_api_key=self.settings.nvidia_api_key,
openrouter_api_key=self.settings.openrouter_api_key,
model_name=self.settings.llm_model,
nvidia_model=self.settings.nvidia_model,
)
# 5. Optional Reranker
self.reranker = None
if self.settings.use_reranker:
try:
self.reranker = CrossEncoderReranker(model_name=self.settings.reranker_model)
except Exception as e:
logger.warning("Could not initialize reranker: %s", e)
# 6. Optional Query Rewriter
self.query_rewriter = None
if self.settings.enable_query_rewriting:
self.query_rewriter = QueryRewriter(llm_service=self.llm_service)
# 7. RAG Pipeline
self.pipeline = RAGPipeline(
retriever=self.retriever,
llm_service=self.llm_service,
reranker=self.reranker,
query_rewriter=self.query_rewriter,
)
def _init_vector_store(self) -> BaseVectorStore:
store_type = self.settings.vector_store.lower()
store_dir = self.settings.get_vectorstore_path()
if store_type == "chroma":
logger.info("Initializing Chroma vector store at '%s'", store_dir)
return ChromaVectorStore(persist_directory=store_dir)
else:
logger.info("Initializing FAISS vector store with dimension %d", self.embedding_service.dimension)
faiss_store = FAISSVectorStore(dimension=self.embedding_service.dimension)
# Attempt to restore if persisted files exist
if (store_dir / "index.faiss").exists() and (store_dir / "metadata.json").exists():
try:
faiss_store.load(store_dir)
logger.info("Restored existing FAISS index with %d chunks", faiss_store.total_chunks())
except Exception as e:
logger.warning("Could not restore FAISS index: %s", e)
return faiss_store
def index_pdf_file(self, file_path: Path) -> dict:
"""Processes a PDF file through extraction, cleaning, chunking, and embedding into the vector store."""
# Clean up existing index for the same filename if previously indexed
for doc in self.vector_store.list_documents():
if doc["filename"] == file_path.name:
logger.info("Replacing existing index for '%s' (ID: %s)", file_path.name, doc["document_id"])
self.vector_store.delete_document(doc["document_id"])
pages = self.pdf_loader.load(file_path)
chunks = self.chunker.chunk_documents(pages)
if not chunks:
raise ValueError(f"No valid text chunks generated from '{file_path.name}'.")
# Generate embeddings in batch
texts = [c.text for c in chunks]
embeddings = self.embedding_service.embed_documents(texts)
# Store in vector database
self.vector_store.add_chunks(chunks, embeddings)
# Persist index
try:
self.vector_store.save(self.settings.get_vectorstore_path())
except Exception as e:
logger.warning("Could not auto-save vector store: %s", e)
return {
"document_id": pages[0].document_id,
"filename": file_path.name,
"pages_extracted": len(pages),
"chunks_created": len(chunks),
"total_indexed_chunks": self.vector_store.total_chunks(),
}
def reload_vector_store_if_needed(self) -> None:
"""Reloads FAISS vector store from disk if persisted files exist and in-memory is empty."""
store_dir = self.settings.get_vectorstore_path()
index_file = store_dir / "index.faiss"
meta_file = store_dir / "metadata.json"
if index_file.exists() and meta_file.exists():
try:
self.vector_store.load(store_dir)
logger.info("Synchronized FAISS index from disk (%d chunks)", self.vector_store.total_chunks())
except Exception as e:
logger.warning("Could not sync FAISS index from disk: %s", e)
# Global app state singleton
_APP_STATE: Optional[AppState] = None
def get_app_state() -> AppState:
"""Returns application state singleton."""
global _APP_STATE
if _APP_STATE is None:
_APP_STATE = AppState()
return _APP_STATE
|