| | from __future__ import annotations |
| |
|
| | import logging |
| | from dataclasses import dataclass |
| | from typing import List, Optional |
| |
|
| | from langchain_core.documents import Document |
| | from langchain_huggingface import HuggingFaceEmbeddings |
| | from langchain_text_splitters import RecursiveCharacterTextSplitter |
| | from langchain_community.vectorstores import FAISS |
| |
|
| | logger = logging.getLogger(__name__) |
| |
|
| |
|
| | @dataclass |
| | class RAGConfig: |
| | chunk_size: int = 1000 |
| | chunk_overlap: int = 100 |
| | embedding_model: str = "sentence-transformers/all-MiniLM-L6-v2" |
| | embedding_device: str = "cpu" |
| | top_k: int = 4 |
| | normalize_embeddings: bool = True |
| |
|
| |
|
| | @dataclass |
| | class RetrievedContext: |
| | query: str |
| | chunks: List[Document] |
| | combined_text: str |
| |
|
| | @property |
| | def source_files(self) -> List[str]: |
| | seen, out = set(), [] |
| | for doc in self.chunks: |
| | s = doc.metadata.get("source", "unknown") |
| | if s not in seen: |
| | seen.add(s) |
| | out.append(s) |
| | return out |
| |
|
| | @property |
| | def chunk_count(self) -> int: |
| | return len(self.chunks) |
| |
|
| |
|
| | class RAGEngine: |
| | def __init__(self, config: Optional[RAGConfig] = None) -> None: |
| | self.config = config or RAGConfig() |
| | self._vector_store: Optional[FAISS] = None |
| | self._embeddings: Optional[HuggingFaceEmbeddings] = None |
| | self._splitter = RecursiveCharacterTextSplitter( |
| | chunk_size=self.config.chunk_size, |
| | chunk_overlap=self.config.chunk_overlap, |
| | separators=["\n\n", "\n", ". ", " ", ""], |
| | add_start_index=True, |
| | ) |
| |
|
| | def build_index(self, documents: List[Document]) -> int: |
| | if not documents: |
| | raise ValueError("Cannot build index from an empty document list.") |
| | chunks = self._chunk(documents) |
| | self._vector_store = FAISS.from_documents(chunks, self._get_embeddings()) |
| | return len(chunks) |
| |
|
| | def retrieve(self, query: str, top_k: Optional[int] = None) -> RetrievedContext: |
| | if self._vector_store is None: |
| | raise RuntimeError("Call build_index() before retrieve().") |
| | docs = self._vector_store.similarity_search(query, k=top_k or self.config.top_k) |
| | return RetrievedContext(query=query, chunks=docs, combined_text=self._format_context(docs)) |
| |
|
| | def add_documents(self, documents: List[Document]) -> int: |
| | if self._vector_store is None: |
| | return self.build_index(documents) |
| | chunks = self._chunk(documents) |
| | self._vector_store.add_documents(chunks, embedding=self._get_embeddings()) |
| | return len(chunks) |
| |
|
| | def reset(self) -> None: |
| | self._vector_store = None |
| |
|
| | def _chunk(self, documents: List[Document]) -> List[Document]: |
| | chunks = self._splitter.split_documents(documents) |
| | for idx, chunk in enumerate(chunks): |
| | chunk.metadata["chunk_id"] = idx |
| | return chunks |
| |
|
| | def _get_embeddings(self) -> HuggingFaceEmbeddings: |
| | if self._embeddings is None: |
| | self._embeddings = HuggingFaceEmbeddings( |
| | model_name=self.config.embedding_model, |
| | model_kwargs={"device": self.config.embedding_device}, |
| | encode_kwargs={"normalize_embeddings": self.config.normalize_embeddings}, |
| | ) |
| | return self._embeddings |
| |
|
| | @staticmethod |
| | def _format_context(docs: List[Document]) -> str: |
| | separator = "─" * 60 |
| | parts = [] |
| | for i, doc in enumerate(docs, start=1): |
| | meta = doc.metadata |
| | header = ( |
| | f"[CHUNK {i} | source: {meta.get('source', 'unknown')}, " |
| | f"page: {meta.get('page', 'N/A')}, " |
| | f"chunk_id: {meta.get('chunk_id', 'N/A')}]" |
| | ) |
| | parts.append(f"{header}\n{doc.page_content.strip()}") |
| | return f"\n{separator}\n".join(parts) |
