File size: 7,544 Bytes
4afc8ce 892a464 4afc8ce 892a464 9d5041f 4afc8ce 892a464 4afc8ce 892a464 9d5041f 892a464 9d5041f 892a464 9d5041f 892a464 9d5041f 892a464 9d5041f 892a464 9d5041f 892a464 9d5041f 892a464 9d5041f 892a464 9d5041f 892a464 9d5041f 892a464 9d5041f 892a464 9d5041f 892a464 9d5041f 892a464 9d5041f 892a464 9d5041f 892a464 9d5041f 892a464 9d5041f 892a464 9d5041f 892a464 | 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 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 | """Knowledge base with RAG capabilities"""
import os
from typing import List
from langchain_openai import OpenAIEmbeddings
from langchain_community.vectorstores import FAISS
from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_core.documents import Document
from config import logger_knowledge
class KnowledgeBase:
"""Knowledge base with FAISS vector store for RAG capabilities"""
def __init__(self, pdf_path: str, index_path: str, embedding_model: str = "text-embedding-3-small", top_k: int = 2, recreate_index: bool = False):
"""
Initialize knowledge base with FAISS vector store
Args:
pdf_path: Path to the PDF document
index_path: Path to save/load the FAISS index
embedding_model: OpenAI embedding model to use
top_k: Number of documents to retrieve
recreate_index: Whether to recreate the FAISS index from scratch
"""
self.pdf_path = pdf_path
self.index_path = index_path
self.top_k = top_k
logger_knowledge.info(f"Initializing KnowledgeBase with embedding_model={embedding_model}, top_k={top_k}")
logger_knowledge.debug(f"PDF path: {pdf_path}")
logger_knowledge.debug(f"Index path: {index_path}")
logger_knowledge.info(f"Loading OpenAI embeddings model: {embedding_model}")
self.embeddings = OpenAIEmbeddings(model=embedding_model)
self.vectorstore = self._load_or_create_index(recreate_index)
def _load_or_create_index(self, recreate: bool = False) -> FAISS:
"""Load existing FAISS index or create new one from PDF"""
# If index exists and not recreating, load it
if not recreate and os.path.exists(self.index_path):
logger_knowledge.info(f"Loading existing FAISS index from {self.index_path}")
try:
vectorstore = FAISS.load_local(
self.index_path,
self.embeddings,
allow_dangerous_deserialization=True
)
logger_knowledge.info("FAISS index loaded successfully")
return vectorstore
except Exception as e:
logger_knowledge.error(f"Failed to load FAISS index: {str(e)}")
raise
# Otherwise, create new index
logger_knowledge.info(f"Creating new FAISS index from {self.pdf_path}")
# Remove old index if recreating
if recreate and os.path.exists(self.index_path):
import shutil
try:
shutil.rmtree(self.index_path)
logger_knowledge.info("Removed old index directory")
except Exception as e:
logger_knowledge.warning(f"Could not remove old index: {e}")
# Load PDF document
if not os.path.exists(self.pdf_path):
error_msg = f"PDF file not found: {self.pdf_path}"
logger_knowledge.error(error_msg)
raise FileNotFoundError(error_msg)
logger_knowledge.info(f"Loading PDF from {self.pdf_path}")
try:
loader = PyPDFLoader(self.pdf_path)
documents = loader.load()
logger_knowledge.info(f"Loaded {len(documents)} pages from PDF")
except Exception as e:
logger_knowledge.error(f"Failed to load PDF: {str(e)}")
raise
# Split documents into chunks using RecursiveCharacterTextSplitter
logger_knowledge.info("Splitting documents into chunks")
text_splitter = RecursiveCharacterTextSplitter(
chunk_size=800, # Optimized for better granularity
chunk_overlap=150, # Reduced proportionally
length_function=len,
separators=["\n\n", "\n", ". ", ", ", " ", ""] # Paragraph > Line > Sentence > Clause > Word
)
chunks = text_splitter.split_documents(documents)
logger_knowledge.info(f"Split into {len(chunks)} chunks")
# Create FAISS index from chunks
logger_knowledge.info("Creating FAISS vector store from chunks")
try:
vectorstore = FAISS.from_documents(chunks, self.embeddings)
logger_knowledge.info("FAISS vector store created successfully")
except Exception as e:
logger_knowledge.error(f"Failed to create FAISS vector store: {str(e)}")
raise
# Save the index
try:
vectorstore.save_local(self.index_path)
logger_knowledge.info(f"Saved FAISS index to {self.index_path}")
except Exception as e:
logger_knowledge.error(f"Failed to save FAISS index: {str(e)}")
raise
return vectorstore
def retrieve_relevant_docs(self, query: str, k: int = None) -> List[Document]:
"""
Retrieve relevant documents for a query
Args:
query: User question
k: Number of documents to retrieve (uses top_k if not specified)
Returns:
List of relevant document chunks
"""
if not self.vectorstore:
logger_knowledge.error("Vector store not initialized!")
return []
k = k or self.top_k
logger_knowledge.debug(f"Retrieving top {k} documents for query")
try:
results = self.vectorstore.similarity_search(query, k=k)
logger_knowledge.info(f"Retrieved {len(results)} documents")
return results
except Exception as e:
logger_knowledge.error(f"Document retrieval failed: {str(e)}")
raise
def retrieve_relevant(self, query: str, k: int = None) -> str:
"""
Retrieve relevant documents as formatted string with metadata
Args:
query: User question
k: Number of documents to retrieve (uses top_k if not specified)
Returns:
Concatenated text from relevant documents with metadata
"""
logger_knowledge.info(f"Retrieving context for query: {query[:50]}..." if len(query) > 50 else f"Retrieving context for query: {query}")
docs = self.retrieve_relevant_docs(query, k)
if not docs:
logger_knowledge.warning("No documents retrieved for query")
return ""
formatted_chunks = []
for i, doc in enumerate(docs, 1):
chunk_text = f"--- Chunk {i} ---"
# Add metadata if available
if doc.metadata:
metadata_str = ", ".join([f"{k}: {v}" for k, v in doc.metadata.items()])
chunk_text += f"\nMetadata: {metadata_str}"
logger_knowledge.debug(f"Chunk {i} metadata: {doc.metadata}")
# Add content
content_preview = doc.page_content[:100] + "..." if len(doc.page_content) > 100 else doc.page_content
logger_knowledge.debug(f"Chunk {i} content preview: {content_preview}")
chunk_text += f"\n\n{doc.page_content}"
formatted_chunks.append(chunk_text)
total_length = sum(len(chunk) for chunk in formatted_chunks)
logger_knowledge.info(f"Formatted {len(formatted_chunks)} chunks, total length: {total_length} characters")
return "\n\n".join(formatted_chunks)
|