| import os |
| from typing import AsyncGenerator, List, Optional |
| from threading import Lock |
|
|
| import chromadb |
| from llama_index.core import Document, Settings, StorageContext, VectorStoreIndex |
| from llama_index.core.chat_engine.types import BaseChatEngine |
| from llama_index.embeddings.fastembed import FastEmbedEmbedding |
| from llama_index.llms.groq import Groq |
| from llama_index.vector_stores.chroma import ChromaVectorStore |
|
|
| from app.config import settings |
| from app.models.schemas import SourceInfo |
|
|
|
|
| _llm_lock = Lock() |
| _llm_initialized = False |
|
|
|
|
| def _ensure_llm() -> None: |
| global _llm_initialized |
| with _llm_lock: |
| if _llm_initialized: |
| return |
| os.environ["GROQ_API_KEY"] = settings.groq_api_key |
| Settings.llm = Groq(model=settings.groq_model, api_key=settings.groq_api_key) |
| Settings.embed_model = FastEmbedEmbedding(model_name=settings.embed_model) |
| _llm_initialized = True |
|
|
|
|
| class RAGService: |
| """Persistent, multi-session RAG service backed by Chroma.""" |
|
|
| def __init__(self) -> None: |
| settings.chroma_dir.mkdir(parents=True, exist_ok=True) |
| self._chroma = chromadb.PersistentClient(path=str(settings.chroma_dir)) |
| self._collection = self._chroma.get_or_create_collection("studyson") |
| self._vector_store = ChromaVectorStore(chroma_collection=self._collection) |
| self._storage = StorageContext.from_defaults(vector_store=self._vector_store) |
| self._index: Optional[VectorStoreIndex] = None |
| self._chat_engines: dict[str, BaseChatEngine] = {} |
| self._indexed_documents: list[str] = self._load_indexed_documents() |
|
|
| def _load_indexed_documents(self) -> list[str]: |
| try: |
| data = self._collection.get(include=["metadatas"]) |
| sources = {m.get("source") for m in (data.get("metadatas") or []) if m} |
| return sorted(s for s in sources if s) |
| except Exception: |
| return [] |
|
|
| def _ensure_index(self) -> VectorStoreIndex: |
| _ensure_llm() |
| if self._index is None: |
| self._index = VectorStoreIndex.from_vector_store( |
| vector_store=self._vector_store, |
| storage_context=self._storage, |
| ) |
| return self._index |
|
|
| def add_document(self, text: str, source_name: str) -> None: |
| index = self._ensure_index() |
| document = Document(text=text, metadata={"source": source_name}) |
| index.insert(document) |
| if source_name not in self._indexed_documents: |
| self._indexed_documents.append(source_name) |
| self._chat_engines.clear() |
|
|
| def _get_chat_engine(self, session_id: str) -> BaseChatEngine: |
| engine = self._chat_engines.get(session_id) |
| if engine is None: |
| index = self._ensure_index() |
| engine = index.as_chat_engine( |
| chat_mode="condense_plus_context", |
| similarity_top_k=settings.similarity_top_k, |
| verbose=False, |
| ) |
| self._chat_engines[session_id] = engine |
| return engine |
|
|
| async def stream_query( |
| self, question: str, session_id: str |
| ) -> AsyncGenerator[str, None]: |
| if not self.has_documents(): |
| raise ValueError("No documents indexed.") |
| engine = self._get_chat_engine(session_id) |
| response = await engine.astream_chat(question) |
| async for token in response.async_response_gen(): |
| yield token |
|
|
| async def query(self, question: str) -> tuple[str, List[SourceInfo]]: |
| if not self.has_documents(): |
| raise ValueError("No documents indexed.") |
| index = self._ensure_index() |
| query_engine = index.as_query_engine(similarity_top_k=settings.similarity_top_k) |
| response = await query_engine.aquery(question) |
| sources: list[SourceInfo] = [] |
| for node in getattr(response, "source_nodes", []) or []: |
| sources.append( |
| SourceInfo( |
| file_name=node.metadata.get("source", "Unknown"), |
| text=node.text[:300], |
| score=getattr(node, "score", None), |
| ) |
| ) |
| return str(response), sources |
|
|
| async def summarize(self, max_length: int = 500) -> str: |
| if not self.has_documents(): |
| raise ValueError("No documents indexed.") |
| index = self._ensure_index() |
| query_engine = index.as_query_engine(similarity_top_k=8) |
| prompt = ( |
| f"Provide a comprehensive summary of all indexed documents in approximately " |
| f"{max_length} words. Cover the main ideas, key arguments, and important details. " |
| f"Use clear paragraphs." |
| ) |
| response = await query_engine.aquery(prompt) |
| return str(response) |
|
|
| def reset_session(self, session_id: str) -> None: |
| self._chat_engines.pop(session_id, None) |
|
|
| def reset_all(self) -> None: |
| try: |
| self._chroma.delete_collection("studyson") |
| except Exception: |
| pass |
| self._collection = self._chroma.get_or_create_collection("studyson") |
| self._vector_store = ChromaVectorStore(chroma_collection=self._collection) |
| self._storage = StorageContext.from_defaults(vector_store=self._vector_store) |
| self._index = None |
| self._chat_engines.clear() |
| self._indexed_documents = [] |
|
|
| def get_indexed_documents(self) -> List[str]: |
| return list(self._indexed_documents) |
|
|
| def has_documents(self) -> bool: |
| try: |
| return self._collection.count() > 0 |
| except Exception: |
| return bool(self._indexed_documents) |
|
|
