Alamgirapi/Professional
0
1import os
2import json
3import faiss
4import numpy as np
5from typing import List, Dict, Optional, Tuple
6import uuid
7from langchain_community.document_loaders import TextLoader, PyPDFLoader
8from langchain.text_splitter import RecursiveCharacterTextSplitter
9from retriever.embeddings import get_embedding_model
10from config import Config
11
12class DocumentStore:
13 """Vector store for document storage and retrieval"""
14
15 def __init__(self, vector_db_path: Optional[str] = None):
16 """Initialize the document store"""
17 self.vector_db_path = vector_db_path or Config.VECTOR_DB_PATH
18 print(f"Using vector DB path: {self.vector_db_path}")
19
20 self.embeddings = get_embedding_model()
21 print("Embedding model loaded")
22
23 self.text_splitter = RecursiveCharacterTextSplitter(
24 chunk_size=1000,
25 chunk_overlap=200
26 )
27
28 # Create directory if it doesn't exist
29 os.makedirs(self.vector_db_path, exist_ok=True)
30
31 # Check if index exists, otherwise create it
32 self.index_path = os.path.join(self.vector_db_path, "faiss_index")
33 self.documents_path = os.path.join(self.vector_db_path, "documents.json")
34
35 print(f"Index path: {self.index_path}")
36 print(f"Documents path: {self.documents_path}")
37
38 # Load or create index
39 if os.path.exists(self.index_path) and os.path.exists(self.documents_path):
40 print("Found existing index and documents, loading...")
41 self.load()
42 else:
43 print("No existing index found, initializing empty one...")
44 # Initialize an empty index
45 self.documents = {}
46 self.document_embeddings = {}
47 self.initialize_index()
48
49 def initialize_index(self):
50 """Initialize an empty FAISS index"""
51 # Get embedding dimension from the model
52 test_embedding = self.embeddings.encode("test")
53 dimension = len(test_embedding)
54
55 # Create empty index
56 self.index = faiss.IndexFlatL2(dimension)
57 self.save()
58
59 def add_text(self, content: str, title: str = "Untitled") -> str:
60 """
61 Add text content to the document store
62
63 Args:
64 content (str): The text content to add
65 title (str): Title for the content
66
67 Returns:
68 str: Document ID
69 """
70 # Generate a unique ID for the document
71 doc_id = str(uuid.uuid4())
72
73 # Split text into chunks
74 chunks = self.text_splitter.split_text(content)
75
76 # Store document metadata
77 self.documents[doc_id] = {
78 "title": title,
79 "chunks": chunks,
80 "type": "text"
81 }
82
83 # Compute and store embeddings for each chunk
84 chunk_embeddings = []
85 for i, chunk in enumerate(chunks):
86 embedding = self.embeddings.encode(chunk)
87 chunk_id = f"{doc_id}_{i}"
88 self.document_embeddings[chunk_id] = {
89 "doc_id": doc_id,
90 "chunk_index": i
91 }
92 chunk_embeddings.append(embedding)
93
94 # Add embeddings to FAISS index
95 if chunk_embeddings:
96 self.index.add(np.array(chunk_embeddings, dtype=np.float32))
97 self.save()
98
99 return doc_id
100
101 def add_document(self, file_path: str) -> str:
102 """
103 Process and add a document file to the store
104
105 Args:
106 file_path (str): Path to the document file
107
108 Returns:
109 str: Document ID
110 """
111 # Determine file type and use appropriate loader
112 if file_path.lower().endswith('.pdf'):
113 loader = PyPDFLoader(file_path)
114 docs = loader.load()
115 elif file_path.lower().endswith('.txt'):
116 loader = TextLoader(file_path)
117 docs = loader.load()
118 else:
119 raise ValueError(f"Unsupported file type: {file_path}")
120
121 # Extract text from documents
122 content = "\n\n".join([doc.page_content for doc in docs])
123 title = os.path.basename(file_path)
124
125 # Add text to document store
126 return self.add_text(content, title)
127
128 def search(self, query: str, top_k: int = 5) -> List[Dict]:
129 """
130 Search for relevant document chunks
131
132 Args:
133 query (str): The search query
134 top_k (int): Number of results to return
135
136 Returns:
137 List[Dict]: List of document chunks with metadata
138 """
139 # Check if there are any documents first
140 if not self.documents:
141 print("No documents in store during search")
142 return []
143
144 # Print debug information
145 print(f"Searching for: {query}")
146 print(f"Document count: {len(self.documents)}")
147 print(f"Document embeddings count: {len(self.document_embeddings)}")
148
149 # Encode the query
150 query_vector = self.embeddings.encode(query)
151 query_vector = np.array([query_vector], dtype=np.float32)
152
153 # Search the index
154 distances, indices = self.index.search(query_vector, top_k)
155 print(f"Search returned {len(indices[0])} results")
156 print(f"Indices: {indices[0]}")
157 print(f"Distances: {distances[0]}")
158
159 results = []
160 for i, idx in enumerate(indices[0]):
161 # Skip invalid indices
162 if idx == -1:
163 continue
164
165 # Skip results with distance above threshold - TEMPORARILY DISABLED FOR DEBUGGING
166 # if distances[0][i] > Config.SIMILARITY_THRESHOLD:
167 # print(f"Skipping result with distance {distances[0][i]} (above threshold {Config.SIMILARITY_THRESHOLD})")
168 # continue
169 print(f"Processing result with distance {distances[0][i]}")
170
171 # Find the corresponding chunk ID
172 chunk_ids = list(self.document_embeddings.keys())
173 if idx >= len(chunk_ids):
174 print(f"Index {idx} out of range for chunk_ids (len: {len(chunk_ids)})")
175 continue
176
177 chunk_id = chunk_ids[idx]
178 chunk_info = self.document_embeddings[chunk_id]
179 doc_id = chunk_info["doc_id"]
180 chunk_index = chunk_info["chunk_index"]
181
182 # Get document content
183 if doc_id not in self.documents:
184 print(f"Document ID {doc_id} not found in documents")
185 continue
186
187 document = self.documents[doc_id]
188 if chunk_index >= len(document["chunks"]):
189 print(f"Chunk index {chunk_index} out of range for document {doc_id}")
190 continue
191
192 chunk_content = document["chunks"][chunk_index]
193
194 print(f"Found relevant chunk: {chunk_content[:50]}...")
195
196 results.append({
197 "content": chunk_content,
198 "title": document["title"],
199 "similarity": float(1 - distances[0][i] / 2), # Normalize similarity score
200 "doc_id": doc_id
201 })
202
203 print(f"Returning {len(results)} results")
204 return results
205 def save(self):
206 """Save the index and documents to disk"""
207 # Save FAISS index
208 faiss.write_index(self.index, self.index_path)
209
210 # Save documents and mappings
211 data = {
212 "documents": self.documents,
213 "document_embeddings": self.document_embeddings
214 }
215 with open(self.documents_path, 'w') as f:
216 json.dump(data, f)
217
218 def load(self):
219 """Load the index and documents from disk"""
220 try:
221 # Load FAISS index
222 self.index = faiss.read_index(self.index_path)
223
224 # Load documents and mappings
225 with open(self.documents_path, 'r') as f:
226 data = json.load(f)
227 self.documents = data.get("documents", {})
228 self.document_embeddings = data.get("document_embeddings", {})
229
230 print(f"Loaded {len(self.documents)} documents and {len(self.document_embeddings)} embeddings")
231
232 # Verify document structure
233 for doc_id, doc in self.documents.items():
234 if "chunks" not in doc:
235 print(f"Warning: Document {doc_id} missing 'chunks' field")
236 elif not doc["chunks"]:
237 print(f"Warning: Document {doc_id} has empty 'chunks' list")
238
239 # Verify embedding-document relationships
240 for chunk_id, chunk_info in self.document_embeddings.items():
241 doc_id = chunk_info.get("doc_id")
242 if doc_id not in self.documents:
243 print(f"Warning: Embedding {chunk_id} refers to non-existent document {doc_id}")
244 continue
245
246 chunk_index = chunk_info.get("chunk_index")
247 if chunk_index is None:
248 print(f"Warning: Embedding {chunk_id} missing 'chunk_index'")
249 continue
250
251 doc = self.documents[doc_id]
252 if "chunks" not in doc or chunk_index >= len(doc["chunks"]):
253 print(f"Warning: Embedding {chunk_id} refers to non-existent chunk {chunk_index} in document {doc_id}")
254
255 except Exception as e:
256 print(f"Error loading document store: {e}")
257 # Initialize empty collections
258 self.documents = {}
259 self.document_embeddings = {}
260 self.initialize_index()
261
262 def rebuild_index(self):
263 """Rebuild the index from all documents"""
264 # Get embedding dimension
265 test_embedding = self.embeddings.encode("test")
266 dimension = len(test_embedding)
267
268 # Create a new index
269 self.index = faiss.IndexFlatL2(dimension)
270
271 # Re-embed and add all chunks
272 all_embeddings = []
273
274 for doc_id, doc_info in self.documents.items():
275 chunks = doc_info.get("chunks", [])
276 for chunk in chunks:
277 embedding = self.embeddings.encode(chunk)
278 all_embeddings.append(embedding)
279
280 if all_embeddings:
281 self.index.add(np.array(all_embeddings, dtype=np.float32))
282
283 self.save()
284
285 def load_from_json(self, json_data):
286 """Load documents from provided JSON data"""
287 self.documents = json_data.get("documents", {})
288 self.document_embeddings = json_data.get("document_embeddings", {})
289
290 # Rebuild the index
291 self.rebuild_index()
292
293 def rebuild_index_from_scratch(self):
294 """Completely rebuild the index from the documents"""
295 print("Rebuilding search index from scratch...")
296
297 # Get embedding dimension
298 test_embedding = self.embeddings.encode("test")
299 dimension = len(test_embedding)
300
301 # Create a new index
302 self.index = faiss.IndexFlatL2(dimension)
303
304 # Track mappings between index positions and document chunks
305 self.document_embeddings = {}
306 current_idx = 0
307
308 # Re-embed and add all chunks
309 all_embeddings = []
310
311 for doc_id, doc_info in self.documents.items():
312 chunks = doc_info.get("chunks", [])
313 print(f"Processing document {doc_id} with {len(chunks)} chunks")
314
315 for i, chunk in enumerate(chunks):
316 embedding = self.embeddings.encode(chunk)
317 all_embeddings.append(embedding)
318
319 # Store mapping
320 chunk_id = f"{doc_id}_{i}"
321 self.document_embeddings[chunk_id] = {
322 "doc_id": doc_id,
323 "chunk_index": i
324 }
325 current_idx += 1
326
327 # Add all embeddings to index at once
328 if all_embeddings:
329 print(f"Adding {len(all_embeddings)} embeddings to index")
330 self.index.add(np.array(all_embeddings, dtype=np.float32))
331 else:
332 print("No embeddings to add to index")
333
334 self.save()
335 print("Index rebuild complete")