diff --git a/implementation/answer_enhanced.py b/implementation/answer_enhanced.py new file mode 100644 index 0000000..27688fd --- /dev/null +++ b/implementation/answer_enhanced.py @@ -0,0 +1,296 @@ +from pathlib import Path +from langchain_openai import ChatOpenAI +from langchain_chroma import Chroma +from langchain_huggingface import HuggingFaceEmbeddings +from langchain_core.messages import SystemMessage, HumanMessage, convert_to_messages +from langchain_core.documents import Document +from sentence_transformers import CrossEncoder +from typing import List +import re +import numpy as np + +from dotenv import load_dotenv + +load_dotenv(override=True) + +MODEL = "gpt-4.1-nano" +DB_NAME = str(Path(__file__).parent.parent / "vector_db") + +# NFL team's proven optimal settings +embeddings = HuggingFaceEmbeddings(model_name="thenlper/gte-small") +RETRIEVAL_K = 20 # NFL team's optimal: retrieve more candidates +RERANK_K = 9 # NFL team's optimal: final number after reranking + +# Enhanced settings for our improvements +USE_QUERY_EXPANSION = True +USE_DYNAMIC_K = True +USE_CONTEXT_COMPRESSION = True + +SYSTEM_PROMPT = """ +You are a knowledgeable, friendly assistant representing the company Insurellm. +You are chatting with a user about Insurellm. +If relevant, use the given context to answer any question. +If you don't know the answer, say so. + +Context (with metadata): +{context} +""" + +vectorstore = Chroma(persist_directory=DB_NAME, embedding_function=embeddings) +retriever = vectorstore.as_retriever(search_kwargs={"k": RETRIEVAL_K}) +llm = ChatOpenAI(temperature=0, model_name=MODEL) + +# Initialize cross-encoder for reranking (NFL team's approach) +_reranker = None +def get_reranker(): + global _reranker + if _reranker is None: + _reranker = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2') + return _reranker + + +def advanced_query_expansion(question: str) -> List[str]: + """ + Enhanced query expansion with domain-specific knowledge and LLM-based expansion. + """ + # Our original domain-specific expansions + expansions = { + "award": ["prize", "recognition", "honor", "achievement", "accolade", "trophy"], + "employee": ["staff", "worker", "team member", "personnel", "colleague"], + "founded": ["established", "created", "started", "launched", "began"], + "contract": ["agreement", "deal", "partnership", "arrangement", "pact"], + "product": ["service", "offering", "solution", "platform", "tool"], + "company": ["organization", "firm", "business", "corporation", "enterprise"], + "year": ["annual", "2023", "2024", "2025", "yearly"], + "location": ["office", "headquarters", "base", "site", "facility"], + "salary": ["pay", "wage", "compensation", "income", "earnings"], + "contract": ["agreement", "deal", "partnership", "arrangement"], + "insurance": ["coverage", "policy", "protection", "assurance"], + "technology": ["tech", "software", "system", "platform", "solution"] + } + + expanded_queries = [question] + question_lower = question.lower() + + # Apply domain-specific expansions + for key, synonyms in expansions.items(): + if key in question_lower: + for synonym in synonyms: + expanded_query = question_lower.replace(key, synonym) + if expanded_query != question_lower and expanded_query not in expanded_queries: + expanded_queries.append(expanded_query) + + # Add question variations for better coverage + variations = [ + question_lower.replace("?", "").strip(), + question_lower.replace("who", "what person").replace("?", "").strip(), + question_lower.replace("when", "what time").replace("?", "").strip(), + question_lower.replace("where", "what location").replace("?", "").strip(), + question_lower.replace("how many", "what number of").replace("?", "").strip(), + ] + + for variation in variations: + if variation and variation not in expanded_queries: + expanded_queries.append(variation) + + return expanded_queries[:5] # Limit to 5 queries to avoid too much noise + + +def calculate_query_complexity(question: str) -> str: + """ + Determine query complexity to adjust retrieval strategy. + """ + question_lower = question.lower() + + # Simple factual questions + if any(word in question_lower for word in ["who", "when", "where", "what is", "how many"]): + if len(question.split()) <= 8: + return "simple" + + # Complex analytical questions + if any(word in question_lower for word in ["compare", "analyze", "explain", "describe", "why", "how"]): + return "complex" + + # Questions requiring multiple pieces of information + if any(word in question_lower for word in ["and", "both", "all", "every", "each"]): + return "multi_fact" + + return "medium" + + +def dynamic_k_selection(question: str, docs: List[Document]) -> int: + """ + Dynamically adjust the number of documents based on query complexity and quality. + """ + if not USE_DYNAMIC_K: + return RERANK_K + + complexity = calculate_query_complexity(question) + + # Calculate document quality scores + query_words = set(re.findall(r'\b\w+\b', question.lower())) + quality_scores = [] + + for doc in docs: + content_words = set(re.findall(r'\b\w+\b', doc.page_content.lower())) + overlap = len(query_words.intersection(content_words)) + quality_scores.append(overlap / len(query_words) if query_words else 0) + + avg_quality = np.mean(quality_scores) if quality_scores else 0 + + # Adjust K based on complexity and quality + if complexity == "simple" and avg_quality > 0.3: + return min(5, len(docs)) # Fewer docs for simple, high-quality matches + elif complexity == "complex": + return min(12, len(docs)) # More docs for complex questions + elif complexity == "multi_fact": + return min(10, len(docs)) # Moderate docs for multi-fact questions + else: + return min(RERANK_K, len(docs)) # Default behavior + + +def compress_context(docs: List[Document], question: str) -> List[Document]: + """ + Compress context by removing redundant information and focusing on relevant parts. + """ + if not USE_CONTEXT_COMPRESSION or len(docs) <= 3: + return docs + + query_words = set(re.findall(r'\b\w+\b', question.lower())) + compressed_docs = [] + + for doc in docs: + # Extract the most relevant sentences + sentences = doc.page_content.split('. ') + relevant_sentences = [] + + for sentence in sentences: + sentence_words = set(re.findall(r'\b\w+\b', sentence.lower())) + overlap = len(query_words.intersection(sentence_words)) + if overlap > 0: + relevant_sentences.append(sentence) + + # If we found relevant sentences, create a compressed version + if relevant_sentences: + compressed_content = '. '.join(relevant_sentences[:3]) # Limit to top 3 sentences + if len(compressed_content) < len(doc.page_content) * 0.7: # Only if significantly shorter + compressed_doc = Document( + page_content=compressed_content, + metadata=doc.metadata + ) + compressed_docs.append(compressed_doc) + else: + compressed_docs.append(doc) + else: + compressed_docs.append(doc) + + return compressed_docs + + +def rerank_documents_cross_encoder(query: str, documents: List[Document]) -> List[Document]: + """ + NFL team's cross-encoder reranking with our enhancements. + """ + reranker = get_reranker() + pairs = [[query, doc.page_content] for doc in documents] + scores = reranker.predict(pairs) + + # Sort by score and return top documents + doc_scores = list(zip(documents, scores)) + doc_scores.sort(key=lambda x: x[1], reverse=True) + + # Apply dynamic K selection + dynamic_k = dynamic_k_selection(query, documents) + top_docs = [doc for doc, score in doc_scores[:dynamic_k]] + + # Apply context compression + compressed_docs = compress_context(top_docs, query) + + return compressed_docs + + +def format_doc_with_metadata(doc: Document, idx: int) -> str: + """ + NFL team's metadata formatting with our enhancements. + """ + meta = doc.metadata + formatted = f"--- Document {idx+1} ---\n" + + # Add structured metadata first + if 'entity_name' in meta: + formatted += f"Entity: {meta['entity_name']}\n" + if 'doc_type' in meta: + formatted += f"Type: {meta['doc_type']}\n" + if 'job_title' in meta: + formatted += f"Job Title: {meta['job_title']}\n" + if 'salary' in meta: + formatted += f"Salary: {meta['salary']}\n" + if 'location' in meta: + formatted += f"Location: {meta['location']}\n" + if 'product_name' in meta: + formatted += f"Product: {meta['product_name']}\n" + if 'client_name' in meta: + formatted += f"Client Name: {meta['client_name']}\n" + if 'contract_number' in meta: + formatted += f"Contract #: {meta['contract_number']}\n" + if 'monthly_payment' in meta: + formatted += f"Payment: {meta['monthly_payment']}\n" + + # Add content with relevance highlighting + content = doc.page_content + query_words = set(re.findall(r'\b\w+\b', meta.get('title', '').lower())) + if query_words: + # Highlight relevant terms in content + for word in query_words: + if len(word) > 3: # Only highlight longer words + content = re.sub(f'\\b{word}\\b', f'**{word}**', content, flags=re.IGNORECASE) + + formatted += f"\nContent:\n{content}\n" + return formatted + + +def fetch_context(question: str) -> list[Document]: + """ + Enhanced retrieval: NFL team's approach + our query expansion + dynamic optimization. + """ + if USE_QUERY_EXPANSION: + # Use our enhanced query expansion + expanded_queries = advanced_query_expansion(question) + + # Retrieve documents for each expanded query + all_documents = [] + for query in expanded_queries: + docs = retriever.invoke(query, k=RETRIEVAL_K) + all_documents.extend(docs) + + # Remove duplicates while preserving order + seen = set() + unique_docs = [] + for doc in all_documents: + doc_id = (doc.page_content, doc.metadata.get('source', '')) + if doc_id not in seen: + seen.add(doc_id) + unique_docs.append(doc) + else: + # Use single query retrieval + unique_docs = retriever.invoke(question, k=RETRIEVAL_K) + + # Use NFL team's cross-encoder reranking with our enhancements + return rerank_documents_cross_encoder(question, unique_docs) + + +def answer_question(question: str, history: list[dict] = []) -> tuple[str, list[Document]]: + """ + Enhanced answer generation with NFL team's approach + our improvements. + """ + docs = fetch_context(question) + + # Format with enhanced metadata for LLM + context = "\n\n".join(format_doc_with_metadata(doc, i) for i, doc in enumerate(docs)) + + system_prompt = SYSTEM_PROMPT.format(context=context) + messages = [SystemMessage(content=system_prompt)] + messages.extend(convert_to_messages(history)) + messages.append(HumanMessage(content=question)) + response = llm.invoke(messages) + return response.content, docs diff --git a/implementation/ingest_enhanced.py b/implementation/ingest_enhanced.py new file mode 100644 index 0000000..4650d3c --- /dev/null +++ b/implementation/ingest_enhanced.py @@ -0,0 +1,230 @@ +import os +import glob +import re +from pathlib import Path +from langchain_community.document_loaders import DirectoryLoader, TextLoader +from langchain_text_splitters import RecursiveCharacterTextSplitter +from langchain_chroma import Chroma +from langchain_huggingface import HuggingFaceEmbeddings + +from dotenv import load_dotenv + +# NFL team's optimal embedding model +MODEL = "thenlper/gte-small" + +DB_NAME = str(Path(__file__).parent.parent / "vector_db") +KNOWLEDGE_BASE = str(Path(__file__).parent.parent / "knowledge-base") + +embeddings = HuggingFaceEmbeddings(model_name=MODEL) + +load_dotenv(override=True) + + +def extract_enhanced_metadata(doc, folder): + """ + Enhanced metadata extraction combining NFL team's approach with our improvements. + """ + metadata = doc.metadata.copy() + doc_type = os.path.basename(folder) + + # Get entity name from filename + filename = Path(doc.metadata['source']).stem + metadata['entity_name'] = filename + metadata['doc_type'] = doc_type + + # Parse document for title + lines = doc.page_content.split('\n') + metadata['title'] = lines[0].replace('#', '').strip() if lines else '' + + # Enhanced entity-specific extraction + if doc_type == 'employees': + # Extract: name, title, salary, location, dob, department + salary_match = re.search(r'\*\*Current Salary:\*\*\s*(\$[\d,]+)', doc.page_content) + if salary_match: + metadata['salary'] = salary_match.group() + + title_match = re.search(r'\*\*Job Title:\*\*\s*(.+)', doc.page_content) + if title_match: + metadata['job_title'] = title_match.group(1).strip() + + location_match = re.search(r'\*\*Location:\*\*\s*(.+)', doc.page_content) + if location_match: + metadata['location'] = location_match.group(1).strip() + + dob_match = re.search(r'\*\*Date of Birth:\*\*\s*(.+)', doc.page_content) + if dob_match: + metadata['dob'] = dob_match.group(1).strip() + + # Extract department/team information + dept_match = re.search(r'\*\*Department:\*\*\s*(.+)', doc.page_content) + if dept_match: + metadata['department'] = dept_match.group(1).strip() + + # Extract years of experience + exp_match = re.search(r'(\d+)\s*years?\s*of\s*experience', doc.page_content, re.IGNORECASE) + if exp_match: + metadata['years_experience'] = exp_match.group(1) + + # Extract skills/technologies + skills_match = re.search(r'\*\*Skills:\*\*\s*(.+)', doc.page_content) + if skills_match: + metadata['skills'] = skills_match.group(1).strip() + + elif doc_type == 'contracts': + # Extract: contract number, client, product, monthly cost, duration + contract_num = re.search(r'\*\*Contract [Number|ID]:\*\*\s*(.+)', doc.page_content) + if contract_num: + metadata['contract_number'] = contract_num.group(1).strip() + + monthly_cost = re.search(r'[monthly payments? of |](\$[\d,]+)[|\sper month]', doc.page_content, re.IGNORECASE) + if monthly_cost: + metadata['monthly_payment'] = monthly_cost.group() + + # Extract contract duration + duration_match = re.search(r'(\d+)\s*(?:month|year)s?\s*(?:contract|term)', doc.page_content, re.IGNORECASE) + if duration_match: + metadata['contract_duration'] = duration_match.group(1) + " " + duration_match.group(2) + + # Extract contract status + status_match = re.search(r'\*\*Status:\*\*\s*(.+)', doc.page_content) + if status_match: + metadata['contract_status'] = status_match.group(1).strip() + + # Extract client and product from filename + if 'Contract with' in filename: + parts = filename.split(' for ') + if len(parts) > 0: + metadata['client_name'] = parts[0].replace('Contract with ', '') + if len(parts) > 1: + metadata['product_name'] = parts[1] + + elif doc_type == 'products': + metadata['product_name'] = filename + + # Extract pricing tiers + pricing_section = re.search(r'## Pricing(.+?)(?=##|\Z)', doc.page_content, re.DOTALL) + if pricing_section: + tier_prices = re.findall(r'\$[\d,]+/month', pricing_section.group(1)) + if tier_prices: + metadata['pricing_info'] = ', '.join(tier_prices) + + # Extract product category + category_match = re.search(r'\*\*Category:\*\*\s*(.+)', doc.page_content) + if category_match: + metadata['product_category'] = category_match.group(1).strip() + + # Extract target audience + audience_match = re.search(r'\*\*Target Audience:\*\*\s*(.+)', doc.page_content) + if audience_match: + metadata['target_audience'] = audience_match.group(1).strip() + + elif doc_type == 'company': + metadata['company_doc'] = filename + + # Extract company values/culture + values_match = re.search(r'\*\*Values:\*\*\s*(.+)', doc.page_content) + if values_match: + metadata['company_values'] = values_match.group(1).strip() + + # Extract company size + size_match = re.search(r'(\d+)\s*employees?', doc.page_content, re.IGNORECASE) + if size_match: + metadata['company_size'] = size_match.group(1) + + # Add content quality indicators + content_length = len(doc.page_content) + metadata['content_length'] = content_length + metadata['has_structured_data'] = bool(re.search(r'\*\*.*\*\*', doc.page_content)) + metadata['has_numbers'] = bool(re.search(r'\d+', doc.page_content)) + metadata['has_dates'] = bool(re.search(r'\b(?:Jan|Feb|Mar|Apr|May|Jun|Jul|Aug|Sep|Oct|Nov|Dec|\d{4})\b', doc.page_content)) + + return metadata + + +def fetch_documents(): + """ + Load documents with enhanced metadata extraction. + """ + folders = glob.glob(str(Path(KNOWLEDGE_BASE) / "*")) + documents = [] + for folder in folders: + loader = DirectoryLoader( + folder, glob="**/*.md", loader_cls=TextLoader, loader_kwargs={"encoding": "utf-8"} + ) + folder_docs = loader.load() + for doc in folder_docs: + # Apply enhanced metadata extraction + doc.metadata = extract_enhanced_metadata(doc, folder) + documents.append(doc) + return documents + + +def create_optimized_chunks(documents): + """ + NFL team's optimal chunking strategy with our enhancements. + """ + # NFL team's proven optimal settings + text_splitter = RecursiveCharacterTextSplitter( + chunk_size=2000, # NFL team's optimal + chunk_overlap=400, # 20% of chunk_size + separators=["\n## ", "\n### ", "\n\n", "\n", " ", ""], # Header-aware + is_separator_regex=False + ) + chunks = text_splitter.split_documents(documents) + + # Add chunk-level metadata for better retrieval + for i, chunk in enumerate(chunks): + chunk.metadata['chunk_id'] = i + chunk.metadata['chunk_length'] = len(chunk.page_content) + + # Add content type indicators + content = chunk.page_content.lower() + chunk.metadata['contains_numbers'] = bool(re.search(r'\d+', content)) + chunk.metadata['contains_dates'] = bool(re.search(r'\b(?:Jan|Feb|Mar|Apr|May|Jun|Jul|Aug|Sep|Oct|Nov|Dec|\d{4})\b', content)) + chunk.metadata['contains_money'] = bool(re.search(r'\$[\d,]+', content)) + chunk.metadata['contains_percentages'] = bool(re.search(r'\d+%', content)) + + # Add semantic indicators + chunk.metadata['is_employee_info'] = any(word in content for word in ['salary', 'job title', 'department', 'employee']) + chunk.metadata['is_contract_info'] = any(word in content for word in ['contract', 'agreement', 'payment', 'client']) + chunk.metadata['is_product_info'] = any(word in content for word in ['product', 'service', 'pricing', 'features']) + chunk.metadata['is_company_info'] = any(word in content for word in ['company', 'founded', 'headquarters', 'mission']) + + return chunks + + +def create_embeddings(chunks): + """ + Create vector store with GTE-small embeddings and enhanced metadata. + """ + if os.path.exists(DB_NAME): + Chroma(persist_directory=DB_NAME, embedding_function=embeddings).delete_collection() + + vectorstore = Chroma.from_documents( + documents=chunks, embedding=embeddings, persist_directory=DB_NAME + ) + + collection = vectorstore._collection + count = collection.count() + + sample_embedding = collection.get(limit=1, include=["embeddings"])["embeddings"][0] + dimensions = len(sample_embedding) + + print(f"There are {count:,} vectors with {dimensions:,} dimensions in the vector store") + print(f"Enhanced metadata fields: {list(chunks[0].metadata.keys()) if chunks else 'None'}") + + return vectorstore + + +if __name__ == "__main__": + print(f"Loading documents from {KNOWLEDGE_BASE}...") + documents = fetch_documents() + print(f"Loaded {len(documents)} documents") + + print("Creating optimized chunks...") + chunks = create_optimized_chunks(documents) + print(f"Created {len(chunks)} chunks") + + print("Creating embeddings and vector store...") + create_embeddings(chunks) + print("Enhanced ingestion complete") diff --git a/test_enhanced.py b/test_enhanced.py new file mode 100644 index 0000000..c804e11 --- /dev/null +++ b/test_enhanced.py @@ -0,0 +1,62 @@ +#!/usr/bin/env python3 +"""Test script for enhanced RAG implementation.""" + +def test_enhanced_implementation(): + print("Testing Enhanced RAG Implementation") + print("=" * 60) + + try: + from implementation.answer_enhanced import answer_question, fetch_context + + # Test cases + test_questions = [ + "Who won the IIOTY award in 2023?", + "How many employees does Insurellm have?", + "What is Insurellm's vision statement?", + "Who founded Insurellm and when?", + "What products does Insurellm offer?" + ] + + for i, question in enumerate(test_questions, 1): + print(f"\n--- Test {i} ---") + print(f"Question: {question}") + + # Test retrieval + docs = fetch_context(question) + print(f"Retrieved {len(docs)} documents") + + # Check document quality + if docs: + print(f"Document types: {[doc.metadata.get('doc_type', 'unknown') for doc in docs[:3]]}") + print(f"Entity names: {[doc.metadata.get('entity_name', 'unknown') for doc in docs[:3]]}") + + # Test full answer + answer, context_docs = answer_question(question) + print(f"Answer: {answer[:100]}...") + print(f"Context docs: {len(context_docs)}") + + # Check for specific improvements + if "IIOTY" in question: + keywords = ['Maxine', 'Thompson', 'IIOTY'] + found_keywords = [] + for doc in context_docs: + content_lower = doc.page_content.lower() + for keyword in keywords: + if keyword.lower() in content_lower: + found_keywords.append(keyword) + print(f"Keywords found: {found_keywords}") + + return True + + except Exception as e: + print(f"Error: {e}") + import traceback + traceback.print_exc() + return False + +if __name__ == "__main__": + success = test_enhanced_implementation() + if success: + print("\n✅ Enhanced implementation test completed successfully!") + else: + print("\n❌ Enhanced implementation test failed!")