Skip to content

Latest commit

 

History

History
320 lines (276 loc) · 13.2 KB

File metadata and controls

320 lines (276 loc) · 13.2 KB

RAG Deep Dive – Full Code Explanation

Date: 2025-09-23 Scope: Current rag/ package (loaders, splitter, embeddings, vectorstore, retriever, chain, wrapper) + how API & UI consume it.


1. Conceptual Model

The Retrieval-Augmented Generation (RAG) layer transforms raw log text into structured, queryable semantic & lexical indices and feeds a language model with the most relevant context for each issue (query). The pipeline stages:

  1. Acquisition (line Documents)
  2. Issue Extraction (filtering for severity)
  3. Context Chunking (windowed text segments)
  4. Embedding (dense vector representations)
  5. Index Construction (FAISS + BM25)
  6. Hybrid Retrieval (ensemble weighted merge)
  7. Prompt Assembly (structured diagnostic template)
  8. Answer Generation (LLM)
  9. Enrichment (citations / timestamps)

2. Data Primitives

LangChain Document

Each Document = { page_content: str, metadata: { source, line_no, severity? } }. Two logical corpora:

  • Issue Documents: individual ERROR/WARN lines.
  • Context Documents: larger chunks built from many adjacent lines (for semantic coverage).

In-Memory Stores

  • ERROR_STORE (api layer) maps issue id → issue record + optional _answer (answer cache + sources).
  • VectorStores instance retains: raw docs, FAISS index, BM25 retriever reference.

3. Module Overview & Responsibilities

Module Responsibility Key Exports
loaders.py Convert raw files to line Documents + severity filtering load_log_files, extract_issue_docs
splitter.py Chunk raw lines into retrieval windows split_context
embeddings.py Provide embedding function/backends get_embeddings (implied)
vectorstore.py Build/augment FAISS + BM25 & persistence VectorStores
retriever.py Compose hybrid ensemble & raw source fetch build_ensemble, fetch_sources
chain.py Build & run RetrievalQA LLM chain build_chain, run_qa
wrapper.py All-in-one programmatic interface RAGWrapper

4. loaders.py – Acquisition & Issue Filtering

Flow

  1. Iterate files line-by-line, constructing a Document per non-empty line.
  2. Detect severity tokens (regex/contains: ERROR, WARN).
  3. Tag severity in metadata for those lines.
  4. extract_issue_docs(docs) returns only documents with severity set.

Design Notes

  • Using line-level Documents preserves source + line number for fine-grained citations.
  • Severity detection intentionally simple for speed; can extend to parse timestamps / log levels via regex groups.
  • Non-blocking: If a file is huge, still O(lines) streaming; memory footprint bounded by keeping lines (optimizable with streaming chunker).

Extension Points

Need Strategy
Additional severities (INFO, DEBUG) Expand predicate & badge mapping
Structured fields (thread id, module) Regex extract & add to metadata
Multiline stack traces Accumulate until blank line → single Document

5. splitter.py – Context Chunking

Purpose

Dense retrieval benefits from semantic windows richer than a single log line. Chunking aggregates adjacent lines up to CHUNK_SIZE chars with overlap to preserve cross-boundary context.

Algorithm (Typical)

Pseudo:

current = [] ; length = 0
for doc in line_docs:
  current.append(doc.page_content)
  length += len(doc.page_content) + 1
  if length >= CHUNK_SIZE:
     emit Document('\n'.join(current), metadata={source: first.source, line_no: first.line_no})
     rewind buffer by overlap characters (approx via slicing concatenated string)

Trade-offs

Parameter Impact
CHUNK_SIZE large Fewer embeddings, broader context per vector; risk of noise
CHUNK_SIZE small More precise retrieval; higher embedding cost
Overlap high Better continuity; more tokens + cost

Enhancements

  • Dynamic splitting on semantic boundaries (timestamps, severity transitions).
  • Skip embedding lines classified as low value (DEBUG noise) to reduce token volume.

6. embeddings.py – Embedding Backend Abstraction

Goals

  • Uniform interface for vectorization.
  • Simple backend switch via env (EMBED_BACKEND = hf or openai).

Typical Implementation

if backend == 'hf':
   return HuggingFaceEmbeddings(model_name)
elif backend == 'openai':
   return OpenAIEmbeddings(model=model_name)

Considerations

Aspect Notes
Normalization Cosine similarity typically uses normalized vectors (FAISS IndexFlatIP or L2)
Batch Size Larger batches improve throughput; watch memory
Model Choice MiniLM for speed; upgrade to all-mpnet-base-v2 for quality

Extensions

  • Add caching (embedding hash -> vector) to skip duplicate line embeddings.
  • Support quantized embeddings for memory reduction.

7. vectorstore.py – Hybrid Index Management

Responsibilities

  • Build and maintain FAISS index over context Documents.
  • Maintain BM25 retriever over the same textual corpus (lexical recall).
  • Allow incremental augmentation (add) merging new docs.
  • Optionally persist to disk for warm restarts.

Core Methods

Method Description
build(issue_docs, context_docs) First-time construction (FAISS + BM25) storing references. Issue docs often stored for potential per-issue queries or future features (clustering).
add(issue_docs, context_docs) Append and rebuild (simple approach). Could be optimized using FAISS add only for vectors.
as_retrievers(k) Returns (vector_retriever, bm25_retriever) configured for top-k retrieval.

Persistence (Conceptual)

  1. Save FAISS index binary.
  2. Dump doc metadata + raw text to JSON.
  3. On startup _try_load() reconstruct embedding dimension (if model consistent) and restore index.

Performance Tuning

Optimization Rationale
Pre-normalize vectors Faster cosine search
Use HNSW index Sub-linear retrieval for large corpora
Shard per file Parallel retrieval then merge

Extension Points

  • Introduce per-severity filtering index (store severity posting lists).
  • Maintain separate indices per tenant / session.
  • Add approximate lexical scorer (e.g., SPLADE) as third signal.

8. retriever.py – Ensemble Composition

Motivation

Semantic vectors and lexical tokens capture complementary signals (e.g., exact error codes vs paraphrased diagnostics). Weighted ensemble improves recall + precision.

Flow

  1. Query both retrievers with same text.
  2. Normalize scores (e.g., min-max or softmax per retriever result set).
  3. Combine: combined = w_vec * vec_score + w_bm25 * bm25_score.
  4. Sort by combined descending, dedupe by doc id / text hash.
  5. Truncate to top-k for downstream.

fetch_sources()

  • Calls ensemble retrieval.
  • Collects minimal citation info: {source, line_no}.
  • Aggregates page_content for potential direct context display.

Tuning Strategy

Scenario Weight Setting
Logs with repeated keywords Lower BM25 weight to reduce redundant lexical hits
Noisy tokens / high variation Increase vector weight
Need exact match guarantee Keep some BM25 weight (>0.2)

Advanced Options

  • Reciprocal rank fusion (RRF) alternative to linear weight.
  • Post-ensemble cross-encoder reranking (quality boost at higher cost).

9. chain.py – LLM RetrievalQA Assembly

Components

  • Prompt Template: Structured headings guide deterministic answer formatting (Diagnosis, Evidence, Remediation Steps, Impact, Uncertainty).
  • RetrievalQA (stuff chain): Concatenates all retrieved docs into a single prompt section.

run_qa()

  1. Invoke chain with {'query': issue_message}.
  2. Chain internally executes retriever -> obtains docs.
  3. Injects docs text into template variables.
  4. Invokes LLM (Chat / Completion) model.
  5. Returns answer; attach original docs as sources with trimmed metadata.

Error Handling

  • If retriever empty: still produce answer (LLM may disclaim) — can enforce fallback text.
  • Add token length guard (truncate docs when exceeding model context limit).

Enhancements

Feature Method
Streaming Replace standard invoke with async streaming & SSE to client
Citation markup Insert inline markers [1][2] referencing doc rank
Confidence score Add self-eval secondary prompt (answer critique)

10. wrapper.py – High-Level Orchestration

Purpose

Allow external Python code to leverage the same ingestion & question answering logic without HTTP. Encapsulates state lifecycle and provides a clean API.

Pattern

rag = RAGWrapper(embedding_backend="hf")
rag.ingest_logs(["/path/app.log"])  # internally loads, extracts issues, chunks
rag.build()                          # if not auto-triggered
ans = rag.query("ERROR connecting to DB")
print(ans.answer, ans.sources)

Internal State

  • Maintains its own VectorStores instance.
  • Tracks whether build is done (status).
  • Provides retrieval-only path for debugging and evaluation.

Extension

  • Multi-corpus mode: maintain dict of corpora keyed by label.
  • Caching: store (query_hash -> answer) to bypass recompute.

11. End-to-End Call Trace (Single Issue)

load_log_files -> extract_issue_docs -> split_context -> VectorStores.build -> _rebuild_chain
resolve -> run_qa -> ensemble retrieval -> LLM -> answer stored -> citations visible

12. Failure Modes & Mitigations

Failure Source Mitigation
Empty retrieval set Sparse logs Fallback message, increase TOP_K
Slow build Large corpus ASYNC_BUILD + progress UI; incremental add batching
Memory pressure Huge embedding matrix Use dimensionality reduction or disk-based index
Inconsistent embeddings after restart Model mismatch Store model id with persisted index & validate before load

13. Performance Optimization Roadmap

Layer Optimization
Embeddings Enable sentence-transformers multi-process, GPU if available
FAISS Switch to HNSW / IVF_FLAT for large scale
BM25 Pre-tokenize & cache term frequencies; optional pruning
Ensemble Implement RRF to reduce parameter tuning overhead
Prompt Dynamic context trimming (ranked tokens until limit)
Batch Resolve Group queries (if similar) and share retrieval set

14. Quality Improvement Hooks

Category Hook
Deduplication Cluster issue messages (MinHash) before indexing
Normalization Preprocess messages (lowercase selective tokens, keep codes)
Enrichment Derive error_type from regex taxonomy for richer prompt context
Feedback Loop Capture accepted answers; fine-tune retrieval weighting

15. Security / Safety Considerations

Concern Vector Layer Mitigation
Prompt Injection in Logs Logs may contain adversarial strings Sanitize / filter suspicious tokens before prompt assembly
PII Leakage Raw logs may contain IDs Add redaction pre-embedding (regex patterns)
Model Misuse Off-policy answers Constrain system prompt to scope (diagnostics only)

16. Testing Suggestions (RAG Focused)

Test Purpose
Retrieval sanity Query known log snippet returns its source top-1
Hybrid parity Vector-only vs hybrid quality delta measured
Chunk boundary Issue near chunk edges still retrieved
Export correctness Ensure resolved flag + timestamps present
Add incremental Second upload doesn't degrade prior recall

17. Common Customizations

Need Change
Adjust recall depth Set TOP_K env; consider differing values per retriever internally
Penalize stale docs Add decay factor to combined score using ingestion timestamp
Multi-tenant Instantiate a VectorStores + QA_CHAIN per tenant id
Severity weighting Re-rank results boosting docs near ERROR lines

18. Minimal Re-Implementation Sketch

(For quick port to another stack)

lines = read_lines(files)
issues = [l for l in lines if is_issue(l)]
chunks = make_chunks(lines)
vecs = embed(chunks)
faiss = build_faiss(vecs)
while query:
  docs_v = search_faiss(query, faiss)
  docs_b = bm25(query, chunks)
  merged = fuse(docs_v, docs_b)
  prompt = make_prompt(query, merged)
  answer = llm(prompt)
  return answer, citations(merged)

19. Quick Reference Table

Stage File Primary Function
Load Lines loaders.py load_log_files
Filter Issues loaders.py extract_issue_docs
Build Context splitter.py split_context
Embed embeddings.py get_embeddings (implied)
Index vectorstore.py build / add
Retrieve retriever.py fetch_sources / ensemble internal
QA Chain chain.py run_qa
Programmatic Wrapper wrapper.py RAGWrapper.query

20. Summary

The rag/ package implements a clear, extensible RAG core: minimal, modular pieces aligned to ingestion, indexing, retrieval fusion, and guided answer generation. Each module has a single responsibility, allowing targeted optimization (performance or quality) without broad refactors. This document serves as a comprehensive guide to understand, extend, and safely optimize the system.

End of File.