Alpha108/GenerativeEngineOptimization
0
1"""
2Vector Chunking and RAG Module
3Handles document chunking, vector embeddings, and RAG question-answering
4"""
5
6import os
7import json
8import numpy as np
9from typing import Dict, Any, List, Optional, Tuple
10from langchain.text_splitter import RecursiveCharacterTextSplitter, CharacterTextSplitter
11from langchain.schema import Document
12from langchain_community.vectorstores import FAISS, Chroma
13from langchain.chains import RetrievalQA, ConversationalRetrievalChain
14from langchain.memory import ConversationBufferMemory
15from langchain.prompts import PromptTemplate
16import tempfile
17import shutil
18
19
20class VectorChunker:
21 """Main class for document chunking and vector operations"""
22
23 def __init__(self, embeddings_model, chunk_size: int = 1000, chunk_overlap: int = 200):
24 self.embeddings = embeddings_model
25 self.chunk_size = chunk_size
26 self.chunk_overlap = chunk_overlap
27 self.setup_text_splitters()
28 self.vector_stores = {} # Cache for vector stores
29
30 def setup_text_splitters(self):
31 """Initialize different text splitting strategies"""
32
33 # Default recursive splitter
34 self.recursive_splitter = RecursiveCharacterTextSplitter(
35 chunk_size=self.chunk_size,
36 chunk_overlap=self.chunk_overlap,
37 length_function=len,
38 separators=["\n\n", "\n", " ", ""]
39 )
40
41 # Character-based splitter
42 self.character_splitter = CharacterTextSplitter(
43 chunk_size=self.chunk_size,
44 chunk_overlap=self.chunk_overlap,
45 separator="\n\n"
46 )
47
48 # Semantic splitter for better context preservation
49 self.semantic_splitter = RecursiveCharacterTextSplitter(
50 chunk_size=800, # Smaller chunks for better semantic coherence
51 chunk_overlap=150,
52 length_function=len,
53 separators=["\n\n", "\n", ". ", " ", ""]
54 )
55
56 def chunk_documents(self, documents: List[Document], strategy: str = "recursive") -> List[Document]:
57 """
58 Chunk documents using specified strategy
59
60 Args:
61 documents (List[Document]): List of documents to chunk
62 strategy (str): Chunking strategy ("recursive", "character", "semantic")
63
64 Returns:
65 List[Document]: List of chunked documents
66 """
67 try:
68 # Choose splitter based on strategy
69 if strategy == "character":
70 splitter = self.character_splitter
71 elif strategy == "semantic":
72 splitter = self.semantic_splitter
73 else:
74 splitter = self.recursive_splitter
75
76 # Split documents
77 chunked_docs = []
78
79 for doc in documents:
80 chunks = splitter.split_documents([doc])
81
82 # Add chunk metadata
83 for i, chunk in enumerate(chunks):
84 chunk.metadata.update({
85 'chunk_index': i,
86 'total_chunks': len(chunks),
87 'chunk_strategy': strategy,
88 'original_source': doc.metadata.get('source', 'unknown'),
89 'chunk_size': len(chunk.page_content),
90 'chunk_word_count': len(chunk.page_content.split())
91 })
92
93 chunked_docs.extend(chunks)
94
95 return chunked_docs
96
97 except Exception as e:
98 raise Exception(f"Document chunking failed: {str(e)}")
99
100 def create_vector_store(self, documents: List[Document], store_type: str = "faiss",
101 persist_directory: Optional[str] = None) -> Any:
102 """
103 Create vector store from documents
104
105 Args:
106 documents (List[Document]): Documents to vectorize
107 store_type (str): Type of vector store ("faiss", "chroma")
108 persist_directory (str): Optional directory to persist the store
109
110 Returns:
111 Vector store instance
112 """
113 try:
114 if not documents:
115 raise ValueError("No documents provided for vector store creation")
116
117 if store_type.lower() == "chroma":
118 if persist_directory:
119 vector_store = Chroma.from_documents(
120 documents=documents,
121 embedding=self.embeddings,
122 persist_directory=persist_directory
123 )
124 vector_store.persist()
125 else:
126 vector_store = Chroma.from_documents(
127 documents=documents,
128 embedding=self.embeddings
129 )
130 else: # Default to FAISS
131 vector_store = FAISS.from_documents(
132 documents=documents,
133 embedding=self.embeddings
134 )
135
136 # Save FAISS index if persist directory provided
137 if persist_directory:
138 os.makedirs(persist_directory, exist_ok=True)
139 vector_store.save_local(persist_directory)
140
141 return vector_store
142
143 except Exception as e:
144 raise Exception(f"Vector store creation failed: {str(e)}")
145
146 def create_qa_chain(self, documents: List[Document], llm, chain_type: str = "stuff") -> RetrievalQA:
147 """
148 Create a Question-Answering chain from documents
149
150 Args:
151 documents (List[Document]): Documents for the knowledge base
152 llm: Language model for answering questions
153 chain_type (str): Type of QA chain ("stuff", "map_reduce", "refine")
154
155 Returns:
156 RetrievalQA: Configured QA chain
157 """
158 try:
159 # Chunk documents
160 chunked_docs = self.chunk_documents(documents, strategy="semantic")
161
162 # Create vector store
163 vector_store = self.create_vector_store(chunked_docs, store_type="faiss")
164
165 # Create retriever
166 retriever = vector_store.as_retriever(
167 search_type="similarity",
168 search_kwargs={"k": 4} # Retrieve top 4 most relevant chunks
169 )
170
171 # Custom prompt for GEO-focused QA
172 qa_prompt_template = """Use the following pieces of context to answer the question at the end.
173If you don't know the answer, just say that you don't know, don't try to make up an answer.
174Focus on providing clear, accurate, and complete answers that would be suitable for AI search engines.
175
176Context:
177{context}
178
179Question: {question}
180
181Answer:"""
182
183 qa_prompt = PromptTemplate(
184 template=qa_prompt_template,
185 input_variables=["context", "question"]
186 )
187
188 # Create QA chain
189 qa_chain = RetrievalQA.from_chain_type(
190 llm=llm,
191 chain_type=chain_type,
192 retriever=retriever,
193 return_source_documents=True,
194 chain_type_kwargs={"prompt": qa_prompt}
195 )
196
197 return qa_chain
198
199 except Exception as e:
200 raise Exception(f"QA chain creation failed: {str(e)}")
201
202 def create_conversational_chain(self, documents: List[Document], llm) -> ConversationalRetrievalChain:
203 """
204 Create a conversational retrieval chain with memory
205
206 Args:
207 documents (List[Document]): Documents for the knowledge base
208 llm: Language model for conversation
209
210 Returns:
211 ConversationalRetrievalChain: Configured conversational chain
212 """
213 try:
214 # Chunk documents
215 chunked_docs = self.chunk_documents(documents, strategy="semantic")
216
217 # Create vector store
218 vector_store = self.create_vector_store(chunked_docs, store_type="faiss")
219
220 # Create retriever
221 retriever = vector_store.as_retriever(
222 search_type="similarity",
223 search_kwargs={"k": 3}
224 )
225
226 # Create memory
227 memory = ConversationBufferMemory(
228 memory_key="chat_history",
229 return_messages=True,
230 output_key="answer"
231 )
232
233 # Custom prompt for conversational QA
234 condense_question_prompt = """Given the following conversation and a follow up question,
235rephrase the follow up question to be a standalone question that can be understood without the chat history.
236
237Chat History:
238{chat_history}
239Follow Up Input: {question}
240Standalone question:"""
241
242 # Create conversational chain
243 conv_chain = ConversationalRetrievalChain.from_llm(
244 llm=llm,
245 retriever=retriever,
246 memory=memory,
247 return_source_documents=True,
248 condense_question_prompt=PromptTemplate.from_template(condense_question_prompt)
249 )
250
251 return conv_chain
252
253 except Exception as e:
254 raise Exception(f"Conversational chain creation failed: {str(e)}")
255
256 def semantic_search(self, query: str, documents: List[Document], top_k: int = 5) -> List[Dict[str, Any]]:
257 """
258 Perform semantic search on documents
259
260 Args:
261 query (str): Search query
262 documents (List[Document]): Documents to search
263 top_k (int): Number of top results to return
264
265 Returns:
266 List[Dict]: Search results with scores
267 """
268 try:
269 # Chunk documents
270 chunked_docs = self.chunk_documents(documents, strategy="semantic")
271
272 # Create vector store
273 vector_store = self.create_vector_store(chunked_docs, store_type="faiss")
274
275 # Perform similarity search with scores
276 results = vector_store.similarity_search_with_score(query, k=top_k)
277
278 # Format results
279 formatted_results = []
280 for doc, score in results:
281 result = {
282 'content': doc.page_content,
283 'metadata': doc.metadata,
284 'similarity_score': float(score),
285 'relevance_rank': len(formatted_results) + 1
286 }
287 formatted_results.append(result)
288
289 return formatted_results
290
291 except Exception as e:
292 raise Exception(f"Semantic search failed: {str(e)}")
293
294 def analyze_document_similarity(self, documents: List[Document]) -> Dict[str, Any]:
295 """
296 Analyze similarity between documents
297
298 Args:
299 documents (List[Document]): Documents to analyze
300
301 Returns:
302 Dict: Similarity analysis results
303 """
304 try:
305 if len(documents) < 2:
306 return {'error': 'Need at least 2 documents for similarity analysis'}
307
308 # Chunk documents
309 chunked_docs = self.chunk_documents(documents, strategy="semantic")
310
311 # Create embeddings for each document
312 doc_embeddings = []
313 doc_metadata = []
314
315 for doc in chunked_docs:
316 # Get embedding for the document
317 embedding = self.embeddings.embed_query(doc.page_content)
318 doc_embeddings.append(embedding)
319 doc_metadata.append({
320 'content_preview': doc.page_content[:200] + "...",
321 'metadata': doc.metadata,
322 'length': len(doc.page_content)
323 })
324
325 # Calculate pairwise similarities
326 similarities = []
327 embeddings_array = np.array(doc_embeddings)
328
329 for i in range(len(embeddings_array)):
330 for j in range(i + 1, len(embeddings_array)):
331 # Calculate cosine similarity
332 similarity = np.dot(embeddings_array[i], embeddings_array[j]) / (
333 np.linalg.norm(embeddings_array[i]) * np.linalg.norm(embeddings_array[j])
334 )
335
336 similarities.append({
337 'doc_1_index': i,
338 'doc_2_index': j,
339 'similarity_score': float(similarity),
340 'doc_1_preview': doc_metadata[i]['content_preview'],
341 'doc_2_preview': doc_metadata[j]['content_preview']
342 })
343
344 # Sort by similarity score
345 similarities.sort(key=lambda x: x['similarity_score'], reverse=True)
346
347 # Calculate statistics
348 similarity_scores = [s['similarity_score'] for s in similarities]
349
350 return {
351 'total_comparisons': len(similarities),
352 'average_similarity': np.mean(similarity_scores),
353 'max_similarity': max(similarity_scores),
354 'min_similarity': min(similarity_scores),
355 'similarity_distribution': {
356 'high_similarity': len([s for s in similarity_scores if s > 0.8]),
357 'medium_similarity': len([s for s in similarity_scores if 0.5 < s <= 0.8]),
358 'low_similarity': len([s for s in similarity_scores if s <= 0.5])
359 },
360 'top_similar_pairs': similarities[:5],
361 'most_dissimilar_pairs': similarities[-3:]
362 }
363
364 except Exception as e:
365 return {'error': f"Similarity analysis failed: {str(e)}"}
366
367 def extract_key_passages(self, documents: List[Document], queries: List[str],
368 passages_per_query: int = 3) -> Dict[str, List[Dict[str, Any]]]:
369 """
370 Extract key passages from documents based on multiple queries
371
372 Args:
373 documents (List[Document]): Documents to search
374 queries (List[str]): List of queries to search for
375 passages_per_query (int): Number of passages to extract per query
376
377 Returns:
378 Dict: Key passages organized by query
379 """
380 try:
381 # Chunk documents
382 chunked_docs = self.chunk_documents(documents, strategy="semantic")
383
384 # Create vector store
385 vector_store = self.create_vector_store(chunked_docs, store_type="faiss")
386
387 key_passages = {}
388
389 for query in queries:
390 # Search for relevant passages
391 results = vector_store.similarity_search_with_score(query, k=passages_per_query)
392
393 passages = []
394 for doc, score in results:
395 passage = {
396 'content': doc.page_content,
397 'relevance_score': float(score),
398 'metadata': doc.metadata,
399 'word_count': len(doc.page_content.split()),
400 'query_match': query
401 }
402 passages.append(passage)
403
404 key_passages[query] = passages
405
406 return key_passages
407
408 except Exception as e:
409 return {'error': f"Key passage extraction failed: {str(e)}"}
410
411 def optimize_chunking_strategy(self, documents: List[Document],
412 test_queries: List[str]) -> Dict[str, Any]:
413 """
414 Test different chunking strategies and recommend the best one
415
416 Args:
417 documents (List[Document]): Documents to test
418 test_queries (List[str]): Queries to test retrieval performance
419
420 Returns:
421 Dict: Optimization results and recommendations
422 """
423 try:
424 strategies = ["recursive", "character", "semantic"]
425 strategy_results = {}
426
427 for strategy in strategies:
428 try:
429 # Test this strategy
430 chunked_docs = self.chunk_documents(documents, strategy=strategy)
431 vector_store = self.create_vector_store(chunked_docs, store_type="faiss")
432
433 # Test retrieval performance
434 retrieval_scores = []
435
436 for query in test_queries:
437 results = vector_store.similarity_search_with_score(query, k=3)
438
439 # Calculate average relevance score
440 if results:
441 avg_score = sum(score for _, score in results) / len(results)
442 retrieval_scores.append(float(avg_score))
443
444 # Calculate strategy metrics
445 avg_retrieval_score = np.mean(retrieval_scores) if retrieval_scores else 0
446 total_chunks = len(chunked_docs)
447 avg_chunk_size = np.mean([len(doc.page_content) for doc in chunked_docs])
448
449 strategy_results[strategy] = {
450 'average_retrieval_score': avg_retrieval_score,
451 'total_chunks': total_chunks,
452 'average_chunk_size': avg_chunk_size,
453 'retrieval_scores': retrieval_scores,
454 'chunk_size_distribution': {
455 'min': min(len(doc.page_content) for doc in chunked_docs),
456 'max': max(len(doc.page_content) for doc in chunked_docs),
457 'std': float(np.std([len(doc.page_content) for doc in chunked_docs]))
458 }
459 }
460
461 except Exception as e:
462 strategy_results[strategy] = {'error': f"Strategy test failed: {str(e)}"}
463
464 # Determine best strategy
465 valid_strategies = {k: v for k, v in strategy_results.items() if 'error' not in v}
466
467 if valid_strategies:
468 best_strategy = max(valid_strategies.keys(),
469 key=lambda k: valid_strategies[k]['average_retrieval_score'])
470
471 recommendation = {
472 'recommended_strategy': best_strategy,
473 'reason': f"Best average retrieval score: {valid_strategies[best_strategy]['average_retrieval_score']:.4f}",
474 'all_results': strategy_results,
475 'performance_summary': {
476 strategy: result.get('average_retrieval_score', 0)
477 for strategy, result in valid_strategies.items()
478 }
479 }
480 else:
481 recommendation = {
482 'recommended_strategy': 'recursive', # Default fallback
483 'reason': 'All strategies failed, using default',
484 'all_results': strategy_results
485 }
486
487 return recommendation
488
489 except Exception as e:
490 return {'error': f"Chunking optimization failed: {str(e)}"}
491
492 def create_document_summary(self, documents: List[Document], llm,
493 summary_type: str = "extractive") -> Dict[str, Any]:
494 """
495 Create document summaries using the chunked content
496
497 Args:
498 documents (List[Document]): Documents to summarize
499 llm: Language model for summarization
500 summary_type (str): Type of summary ("extractive", "abstractive")
501
502 Returns:
503 Dict: Summary results
504 """
505 try:
506 # Chunk documents for better processing
507 chunked_docs = self.chunk_documents(documents, strategy="semantic")
508
509 if summary_type == "extractive":
510 # Extract key sentences/chunks
511 return self._create_extractive_summary(chunked_docs)
512 else:
513 # Generate abstractive summary using LLM
514 return self._create_abstractive_summary(chunked_docs, llm)
515
516 except Exception as e:
517 return {'error': f"Document summarization failed: {str(e)}"}
518
519 def _create_extractive_summary(self, chunked_docs: List[Document]) -> Dict[str, Any]:
520 """Create extractive summary by selecting key chunks"""
521 try:
522 # Simple extractive approach: select chunks with highest semantic density
523 chunk_scores = []
524
525 for doc in chunked_docs:
526 content = doc.page_content
527 # Simple scoring based on content characteristics
528 word_count = len(content.split())
529 sentence_count = len([s for s in content.split('.') if s.strip()])
530
531 # Score based on information density
532 density_score = word_count / max(sentence_count, 1)
533
534 # Bonus for chunks with questions, definitions, or lists
535 structure_bonus = 0
536 if '?' in content:
537 structure_bonus += 1
538 if any(word in content.lower() for word in ['define', 'definition', 'means', 'refers to']):
539 structure_bonus += 2
540 if content.count('\n•') > 0 or content.count('1.') > 0:
541 structure_bonus += 1
542
543 total_score = density_score + structure_bonus
544 chunk_scores.append((doc, total_score))
545
546 # Sort by score and select top chunks for summary
547 chunk_scores.sort(key=lambda x: x[1], reverse=True)
548 top_chunks = chunk_scores[:min(5, len(chunk_scores))]
549
550 summary_content = []
551 for doc, score in top_chunks:
552 summary_content.append({
553 'content': doc.page_content,
554 'score': score,
555 'metadata': doc.metadata
556 })
557
558 return {
559 'summary_type': 'extractive',
560 'key_chunks': summary_content,
561 'total_chunks_analyzed': len(chunked_docs),
562 'chunks_selected': len(top_chunks)
563 }
564
565 except Exception as e:
566 return {'error': f"Extractive summary failed: {str(e)}"}
567
568 def _create_abstractive_summary(self, chunked_docs: List[Document], llm) -> Dict[str, Any]:
569 """Create abstractive summary using language model"""
570 try:
571 # Combine content from top chunks
572 combined_content = "\n\n".join([doc.page_content for doc in chunked_docs[:10]])
573
574 summary_prompt = f"""Please provide a comprehensive summary of the following content.
575Focus on the main topics, key insights, and important details that would be valuable for AI search engines.
576
577Content:
578{combined_content[:5000]}
579
580Summary:"""
581
582 from langchain.prompts import ChatPromptTemplate
583
584 prompt_template = ChatPromptTemplate.from_messages([
585 ("system", "You are a professional content summarizer. Create clear, informative summaries."),
586 ("user", summary_prompt)
587 ])
588
589 chain = prompt_template | llm
590 result = chain.invoke({})
591
592 summary_text = result.content if hasattr(result, 'content') else str(result)
593
594 return {
595 'summary_type': 'abstractive',
596 'summary': summary_text,
597 'source_chunks': len(chunked_docs),
598 'content_length_processed': len(combined_content)
599 }
600
601 except Exception as e:
602 return {'error': f"Abstractive summary failed: {str(e)}"}
603
604 def save_vector_store(self, vector_store, directory_path: str, store_type: str = "faiss") -> bool:
605 """
606 Save vector store to disk
607
608 Args:
609 vector_store: Vector store instance to save
610 directory_path (str): Directory to save the store
611 store_type (str): Type of vector store
612
613 Returns:
614 bool: Success status
615 """
616 try:
617 os.makedirs(directory_path, exist_ok=True)
618
619 if store_type.lower() == "faiss":
620 vector_store.save_local(directory_path)
621 elif store_type.lower() == "chroma":
622 # Chroma stores are typically persisted during creation
623 pass
624
625 return True
626
627 except Exception as e:
628 print(f"Failed to save vector store: {str(e)}")
629 return False
630
631 def load_vector_store(self, directory_path: str, store_type: str = "faiss"):
632 """
633 Load vector store from disk
634
635 Args:
636 directory_path (str): Directory containing the saved store
637 store_type (str): Type of vector store
638
639 Returns:
640 Vector store instance or None if failed
641 """
642 try:
643 if not os.path.exists(directory_path):
644 return None
645
646 if store_type.lower() == "faiss":
647 vector_store = FAISS.load_local(
648 directory_path,
649 self.embeddings,
650 allow_dangerous_deserialization=True
651 )
652 return vector_store
653 elif store_type.lower() == "chroma":
654 vector_store = Chroma(
655 persist_directory=directory_path,
656 embedding_function=self.embeddings
657 )
658 return vector_store
659
660 return None
661
662 except Exception as e:
663 print(f"Failed to load vector store: {str(e)}")
664 return None
665
666 def get_chunking_stats(self, documents: List[Document], strategy: str = "recursive") -> Dict[str, Any]:
667 """
668 Get detailed statistics about document chunking
669
670 Args:
671 documents (List[Document]): Documents to analyze
672 strategy (str): Chunking strategy to use
673
674 Returns:
675 Dict: Detailed chunking statistics
676 """
677 try:
678 # Chunk documents
679 chunked_docs = self.chunk_documents(documents, strategy=strategy)
680
681 # Calculate statistics
682 chunk_sizes = [len(doc.page_content) for doc in chunked_docs]
683 word_counts = [len(doc.page_content.split()) for doc in chunked_docs]
684
685 stats = {
686 'strategy_used': strategy,
687 'original_documents': len(documents),
688 'total_chunks': len(chunked_docs),
689 'chunk_size_stats': {
690 'min': min(chunk_sizes) if chunk_sizes else 0,
691 'max': max(chunk_sizes) if chunk_sizes else 0,
692 'mean': np.mean(chunk_sizes) if chunk_sizes else 0,
693 'median': np.median(chunk_sizes) if chunk_sizes else 0,
694 'std': np.std(chunk_sizes) if chunk_sizes else 0
695 },
696 'word_count_stats': {
697 'min': min(word_counts) if word_counts else 0,
698 'max': max(word_counts) if word_counts else 0,
699 'mean': np.mean(word_counts) if word_counts else 0,
700 'median': np.median(word_counts) if word_counts else 0,
701 'std': np.std(word_counts) if word_counts else 0
702 },
703 'chunk_distribution': {
704 'very_small': len([s for s in chunk_sizes if s < 200]),
705 'small': len([s for s in chunk_sizes if 200 <= s < 500]),
706 'medium': len([s for s in chunk_sizes if 500 <= s < 1000]),
707 'large': len([s for s in chunk_sizes if 1000 <= s < 2000]),
708 'very_large': len([s for s in chunk_sizes if s >= 2000])
709 },
710 'overlap_efficiency': self._calculate_overlap_efficiency(chunked_docs),
711 'content_coverage': self._calculate_content_coverage(documents, chunked_docs)
712 }
713
714 return stats
715
716 except Exception as e:
717 return {'error': f"Chunking statistics failed: {str(e)}"}
718
719 def _calculate_overlap_efficiency(self, chunked_docs: List[Document]) -> float:
720 """Calculate efficiency of chunk overlaps"""
721 try:
722 if len(chunked_docs) < 2:
723 return 1.0
724
725 total_content_length = sum(len(doc.page_content) for doc in chunked_docs)
726 unique_content = set()
727
728 # Rough estimate of content uniqueness
729 for doc in chunked_docs:
730 words = doc.page_content.split()
731 for i in range(0, len(words), 10): # Sample every 10th word
732 unique_content.add(' '.join(words[i:i+10]))
733
734 # Efficiency as ratio of unique content to total content
735 efficiency = len(unique_content) * 10 / total_content_length if total_content_length > 0 else 0
736 return min(efficiency, 1.0)
737
738 except Exception:
739 return 0.5 # Default neutral efficiency
740
741 def _calculate_content_coverage(self, original_docs: List[Document],
742 chunked_docs: List[Document]) -> float:
743 """Calculate how well chunks cover original content"""
744 try:
745 original_content = ' '.join([doc.page_content for doc in original_docs])
746 chunked_content = ' '.join([doc.page_content for doc in chunked_docs])
747
748 # Simple coverage metric based on length
749 coverage = len(chunked_content) / len(original_content) if original_content else 0
750 return min(coverage, 1.0)
751
752 except Exception:
753 return 0.0
754
755
756class ChunkingOptimizer:
757 """Helper class for optimizing chunking parameters"""
758
759 def __init__(self, embeddings_model):
760 self.embeddings = embeddings_model
761
762 def optimize_chunk_size(self, documents: List[Document], test_queries: List[str],
763 size_range: Tuple[int, int] = (200, 2000),
764 step_size: int = 200) -> Dict[str, Any]:
765 """
766 Find optimal chunk size for given documents and queries
767
768 Args:
769 documents (List[Document]): Documents to test
770 test_queries (List[str]): Queries for testing retrieval
771 size_range (Tuple[int, int]): Range of chunk sizes to test
772 step_size (int): Step size for testing
773
774 Returns:
775 Dict: Optimization results with recommended chunk size
776 """
777 try:
778 results = {}
779 min_size, max_size = size_range
780
781 for chunk_size in range(min_size, max_size + 1, step_size):
782 # Test this chunk size
783 chunker = VectorChunker(self.embeddings, chunk_size=chunk_size)
784
785 try:
786 chunked_docs = chunker.chunk_documents(documents)
787 vector_store = chunker.create_vector_store(chunked_docs)
788
789 # Test retrieval performance
790 retrieval_scores = []
791 for query in test_queries:
792 search_results = vector_store.similarity_search_with_score(query, k=3)
793 if search_results:
794 avg_score = sum(score for _, score in search_results) / len(search_results)
795 retrieval_scores.append(float(avg_score))
796
797 avg_performance = np.mean(retrieval_scores) if retrieval_scores else 0
798
799 results[chunk_size] = {
800 'average_retrieval_score': avg_performance,
801 'total_chunks': len(chunked_docs),
802 'retrieval_scores': retrieval_scores
803 }
804
805 except Exception as e:
806 results[chunk_size] = {'error': str(e)}
807
808 # Find optimal chunk size
809 valid_results = {k: v for k, v in results.items() if 'error' not in v}
810
811 if valid_results:
812 optimal_size = max(valid_results.keys(),
813 key=lambda k: valid_results[k]['average_retrieval_score'])
814
815 return {
816 'optimal_chunk_size': optimal_size,
817 'optimal_performance': valid_results[optimal_size]['average_retrieval_score'],
818 'all_results': results,
819 'performance_trend': self._analyze_performance_trend(valid_results),
820 'recommendation': f"Use chunk size {optimal_size} for best retrieval performance"
821 }
822 else:
823 return {
824 'error': 'No valid chunk sizes could be tested',
825 'all_results': results
826 }
827
828 except Exception as e:
829 return {'error': f"Chunk size optimization failed: {str(e)}"}
830
831 def _analyze_performance_trend(self, results: Dict[int, Dict[str, Any]]) -> Dict[str, Any]:
832 """Analyze performance trend across different chunk sizes"""
833 try:
834 sizes = sorted(results.keys())
835 performances = [results[size]['average_retrieval_score'] for size in sizes]
836
837 # Find trend direction
838 if len(performances) >= 2:
839 trend_direction = "increasing" if performances[-1] > performances[0] else "decreasing"
840 peak_performance = max(performances)
841 peak_size = sizes[performances.index(peak_performance)]
842
843 return {
844 'trend_direction': trend_direction,
845 'peak_performance': peak_performance,
846 'peak_size': peak_size,
847 'performance_range': max(performances) - min(performances),
848 'stable_performance': max(performances) - min(performances) < 0.1
849 }
850 else:
851 return {'error': 'Insufficient data for trend analysis'}
852
853 except Exception:
854 return {'error': 'Trend analysis failed'}
855
856
857class RAGPipeline:
858 """Complete RAG pipeline for document question-answering"""
859
860 def __init__(self, embeddings_model, llm):
861 self.embeddings = embeddings_model
862 self.llm = llm
863 self.chunker = VectorChunker(embeddings_model)
864 self.vector_stores = {}
865 self.qa_chains = {}
866
867 def create_pipeline(self, documents: List[Document], pipeline_id: str,
868 chunking_strategy: str = "semantic") -> Dict[str, Any]:
869 """
870 Create a complete RAG pipeline for documents
871
872 Args:
873 documents (List[Document]): Documents to process
874 pipeline_id (str): Unique identifier for this pipeline
875 chunking_strategy (str): Strategy for document chunking
876
877 Returns:
878 Dict: Pipeline creation results
879 """
880 try:
881 # Step 1: Chunk documents
882 chunked_docs = self.chunker.chunk_documents(documents, strategy=chunking_strategy)
883
884 # Step 2: Create vector store
885 vector_store = self.chunker.create_vector_store(chunked_docs, store_type="faiss")
886
887 # Step 3: Create QA chain
888 qa_chain = self.chunker.create_qa_chain(documents, self.llm)
889
890 # Store pipeline components
891 self.vector_stores[pipeline_id] = vector_store
892 self.qa_chains[pipeline_id] = qa_chain
893
894 # Pipeline statistics
895 stats = {
896 'pipeline_id': pipeline_id,
897 'documents_processed': len(documents),
898 'chunks_created': len(chunked_docs),
899 'chunking_strategy': chunking_strategy,
900 'vector_store_type': 'faiss',
901 'embedding_model': str(self.embeddings),
902 'created_at': self._get_timestamp()
903 }
904
905 return {
906 'success': True,
907 'pipeline_stats': stats,
908 'chunking_info': self.chunker.get_chunking_stats(documents, chunking_strategy)
909 }
910
911 except Exception as e:
912 return {'error': f"Pipeline creation failed: {str(e)}"}
913
914 def query_pipeline(self, pipeline_id: str, query: str,
915 return_sources: bool = True) -> Dict[str, Any]:
916 """
917 Query a created RAG pipeline
918
919 Args:
920 pipeline_id (str): ID of the pipeline to query
921 query (str): Question to ask
922 return_sources (bool): Whether to return source documents
923
924 Returns:
925 Dict: Query results with answer and sources
926 """
927 try:
928 if pipeline_id not in self.qa_chains:
929 return {'error': f"Pipeline '{pipeline_id}' not found"}
930
931 qa_chain = self.qa_chains[pipeline_id]
932
933 # Execute query
934 result = qa_chain({"query": query})
935
936 # Format response
937 response = {
938 'query': query,
939 'answer': result.get('result', 'No answer generated'),
940 'pipeline_id': pipeline_id,
941 'query_timestamp': self._get_timestamp()
942 }
943
944 # Add source documents if requested
945 if return_sources and 'source_documents' in result:
946 sources = []
947 for i, doc in enumerate(result['source_documents']):
948 source = {
949 'source_index': i,
950 'content': doc.page_content,
951 'metadata': doc.metadata,
952 'relevance_rank': i + 1
953 }
954 sources.append(source)
955
956 response['sources'] = sources
957 response['num_sources'] = len(sources)
958
959 return response
960
961 except Exception as e:
962 return {'error': f"Pipeline query failed: {str(e)}"}
963
964 def batch_query_pipeline(self, pipeline_id: str, queries: List[str]) -> List[Dict[str, Any]]:
965 """
966 Execute multiple queries on a pipeline
967
968 Args:
969 pipeline_id (str): ID of the pipeline to query
970 queries (List[str]): List of questions to ask
971
972 Returns:
973 List[Dict]: List of query results
974 """
975 results = []
976
977 for i, query in enumerate(queries):
978 try:
979 result = self.query_pipeline(pipeline_id, query, return_sources=False)
980 result['batch_index'] = i
981 results.append(result)
982
983 except Exception as e:
984 results.append({
985 'batch_index': i,
986 'query': query,
987 'error': f"Batch query failed: {str(e)}"
988 })
989
990 return results
991
992 def evaluate_pipeline(self, pipeline_id: str, test_queries: List[str],
993 expected_answers: List[str] = None) -> Dict[str, Any]:
994 """
995 Evaluate pipeline performance on test queries
996
997 Args:
998 pipeline_id (str): ID of the pipeline to evaluate
999 test_queries (List[str]): Test questions
1000 expected_answers (List[str]): Optional expected answers for comparison
1001
1002 Returns:
1003 Dict: Evaluation results
1004 """
1005 try:
1006 if pipeline_id not in self.qa_chains:
1007 return {'error': f"Pipeline '{pipeline_id}' not found"}
1008
1009 evaluation_results = []
1010 response_times = []
1011
1012 for i, query in enumerate(test_queries):
1013 import time
1014 start_time = time.time()
1015
1016 # Execute query
1017 result = self.query_pipeline(pipeline_id, query, return_sources=True)
1018
1019 end_time = time.time()
1020 response_time = end_time - start_time
1021 response_times.append(response_time)
1022
1023 # Evaluate result
1024 eval_result = {
1025 'query_index': i,
1026 'query': query,
1027 'answer_generated': not result.get('error'),
1028 'response_time': response_time,
1029 'answer_length': len(result.get('answer', '')),
1030 'sources_returned': result.get('num_sources', 0)
1031 }
1032
1033 # If expected answer provided, calculate similarity
1034 if expected_answers and i < len(expected_answers):
1035 expected = expected_answers[i]
1036 generated = result.get('answer', '')
1037
1038 # Simple similarity metric
1039 similarity = self._calculate_answer_similarity(expected, generated)
1040 eval_result['answer_similarity'] = similarity
1041 eval_result['expected_answer'] = expected
1042
1043 evaluation_results.append(eval_result)
1044
1045 # Calculate aggregate metrics
1046 successful_queries = len([r for r in evaluation_results if r['answer_generated']])
1047 avg_response_time = np.mean(response_times) if response_times else 0
1048
1049 if expected_answers:
1050 similarities = [r.get('answer_similarity', 0) for r in evaluation_results
1051 if 'answer_similarity' in r]
1052 avg_similarity = np.mean(similarities) if similarities else 0
1053 else:
1054 avg_similarity = None
1055
1056 return {
1057 'pipeline_id': pipeline_id,
1058 'total_queries': len(test_queries),
1059 'successful_queries': successful_queries,
1060 'success_rate': successful_queries / len(test_queries) if test_queries else 0,
1061 'average_response_time': avg_response_time,
1062 'average_answer_similarity': avg_similarity,
1063 'detailed_results': evaluation_results,
1064 'evaluation_timestamp': self._get_timestamp()
1065 }
1066
1067 except Exception as e:
1068 return {'error': f"Pipeline evaluation failed: {str(e)}"}
1069
1070 def _calculate_answer_similarity(self, expected: str, generated: str) -> float:
1071 """Calculate similarity between expected and generated answers"""
1072 try:
1073 # Simple word overlap similarity
1074 expected_words = set(expected.lower().split())
1075 generated_words = set(generated.lower().split())
1076
1077 if not expected_words and not generated_words:
1078 return 1.0
1079
1080 intersection = expected_words.intersection(generated_words)
1081 union = expected_words.union(generated_words)
1082
1083 return len(intersection) / len(union) if union else 0.0
1084
1085 except Exception:
1086 return 0.0
1087
1088 def get_pipeline_info(self, pipeline_id: str) -> Dict[str, Any]:
1089 """Get information about a specific pipeline"""
1090 try:
1091 if pipeline_id not in self.qa_chains:
1092 return {'error': f"Pipeline '{pipeline_id}' not found"}
1093
1094 # Get vector store info
1095 vector_store = self.vector_stores.get(pipeline_id)
1096 if vector_store:
1097 try:
1098 # Try to get vector store statistics
1099 total_vectors = vector_store.index.ntotal if hasattr(vector_store, 'index') else 'unknown'
1100 except:
1101 total_vectors = 'unknown'
1102 else:
1103 total_vectors = 'unknown'
1104
1105 return {
1106 'pipeline_id': pipeline_id,
1107 'has_qa_chain': pipeline_id in self.qa_chains,
1108 'has_vector_store': pipeline_id in self.vector_stores,
1109 'total_vectors': total_vectors,
1110 'embedding_model': str(self.embeddings),
1111 'llm_model': str(self.llm)
1112 }
1113
1114 except Exception as e:
1115 return {'error': f"Failed to get pipeline info: {str(e)}"}
1116
1117 def list_pipelines(self) -> Dict[str, Any]:
1118 """List all created pipelines"""
1119 return {
1120 'total_pipelines': len(self.qa_chains),
1121 'pipeline_ids': list(self.qa_chains.keys()),
1122 'vector_stores': list(self.vector_stores.keys())
1123 }
1124
1125 def delete_pipeline(self, pipeline_id: str) -> Dict[str, Any]:
1126 """Delete a pipeline and free resources"""
1127 try:
1128 deleted_components = []
1129
1130 if pipeline_id in self.qa_chains:
1131 del self.qa_chains[pipeline_id]
1132 deleted_components.append('qa_chain')
1133
1134 if pipeline_id in self.vector_stores:
1135 del self.vector_stores[pipeline_id]
1136 deleted_components.append('vector_store')
1137
1138 if deleted_components:
1139 return {
1140 'success': True,
1141 'pipeline_id': pipeline_id,
1142 'deleted_components': deleted_components
1143 }
1144 else:
1145 return {'error': f"Pipeline '{pipeline_id}' not found"}
1146
1147 except Exception as e:
1148 return {'error': f"Pipeline deletion failed: {str(e)}"}
1149
1150 def export_pipeline_config(self, pipeline_id: str) -> Dict[str, Any]:
1151 """Export pipeline configuration for recreation"""
1152 try:
1153 if pipeline_id not in self.qa_chains:
1154 return {'error': f"Pipeline '{pipeline_id}' not found"}
1155
1156 config = {
1157 'pipeline_id': pipeline_id,
1158 'embedding_model_name': getattr(self.embeddings, 'model_name', 'unknown'),
1159 'llm_model_name': getattr(self.llm, 'model_name', 'unknown'),
1160 'chunker_config': {
1161 'chunk_size': self.chunker.chunk_size,
1162 'chunk_overlap': self.chunker.chunk_overlap
1163 },
1164 'export_timestamp': self._get_timestamp(),
1165 'vector_store_type': 'faiss'
1166 }
1167
1168 return config
1169
1170 except Exception as e:
1171 return {'error': f"Pipeline export failed: {str(e)}"}
1172
1173 def _get_timestamp(self) -> str:
1174 """Get current timestamp"""
1175 from datetime import datetime
1176 return datetime.now().strftime('%Y-%m-%d %H:%M:%S')
1177
1178
1179# Utility functions for the module
1180
1181def optimize_rag_pipeline(documents: List[Document], embeddings_model, llm,
1182 test_queries: List[str]) -> Dict[str, Any]:
1183 """
1184 Optimize RAG pipeline configuration for given documents and queries
1185
1186 Args:
1187 documents (List[Document]): Documents to optimize for
1188 embeddings_model: Embedding model to use
1189 llm: Language model to use
1190 test_queries (List[str]): Test queries for optimization
1191
1192 Returns:
1193 Dict: Optimization recommendations
1194 """
1195 try:
1196 # Test different chunking strategies
1197 chunker = VectorChunker(embeddings_model)
1198 chunking_results = chunker.optimize_chunking_strategy(documents, test_queries)
1199
1200 # Test different chunk sizes