Tutorial D7: RAG Foundations (LlamaIndex)

from dataclasses import dataclass

@dataclass
class Doc:
    source: str  # Document identifier (e.g., vendor manual, SOP, PLC code)
    text: str    # Full document text

# Hypothetical cost calculation for prompt stuffing
# Example: 2 vendor PDFs (200 pages each) ≈ 400k-1M tokens
# At $0.15/1M input tokens (gpt-4o-mini as of 2025):
#   400k tokens × $0.15/1M = $0.06 per query
#   1M tokens × $0.15/1M = $0.15 per query
#
# Problem: cost grows with doc set size, not query complexity
# 1,000 queries/day × $0.15 = $150/day = $4,500/month
#
# Additional issues:
# - No stable citations (can't reference specific lines)
# - Answers drift as docs change (no version hash)
# - Context window limits (can't fit 10 manuals)

import math
import textwrap

docs = [
    Doc(
        source="vendor_manual_sinamics_g120.txt",
        text=textwrap.dedent("""
        SINAMICS G120 — Fault Codes (Excerpt)
        F30005: DC link overvoltage.
        Possible causes:
          - Regenerative energy during deceleration (ramp too fast).
          - Line voltage too high.
          - Braking resistor missing or undersized.
        Recommended actions:
          - Increase deceleration ramp time (P1121).
          - Enable/size braking resistor; verify wiring.
          - Check incoming line voltage.
        Notes: Clear fault after cause removed. Document ID: G120-FLT-EX-01
        """).strip()
    ),
    Doc(
        source="line_sop_packaging_a.txt",
        text=textwrap.dedent("""
        Packaging Line A — Troubleshooting SOP (Excerpt)
        Symptom: Conveyor stops after ~3 seconds, HMI shows "Photoeye Timeout".
        Checklist:
          1) Verify photoeye PE-17 alignment and clean lens.
          2) Check PE-17 input in PLC tag table; confirm it toggles when blocked/unblocked.
          3) Inspect conveyor motor overload relay OL-2; reset if tripped.
          4) Review last 10 alarms in SCADA for recurring stop causes.
        Safety: Lockout/Tagout required before mechanical intervention.
        SOP ID: SOP-PL-A-07
        """).strip()
    ),
    Doc(
        source="plc_code_conveyor_st.txt",
        text=textwrap.dedent("""
        // Structured Text (Excerpt) — Conveyor Control
        IF StartButton AND NOT EmergencyStop THEN
            MotorRunning := TRUE;
            StartTimer(IN:=TRUE, PT:=T#3S); // Start grace period
        END_IF

        // Fault: photoeye should clear within 3 seconds of motor start
        IF MotorRunning AND PhotoEye_Blocked THEN
            IF StartTimer.Q THEN
                Fault_PhotoEye_Timeout := TRUE;
                MotorRunning := FALSE;
            END_IF
        END_IF

        // Drive fault mapping
        IF DriveFault_F30005 THEN
            AlarmText := 'DC link overvoltage';
        END_IF
        """).strip()
    ),
    Doc(
        source="change_log_2025w41.txt",
        text=textwrap.dedent("""
        Change Log — Packaging Line A (Week 41)
        - Updated VFD decel ramp: P1121 from 1.0s to 0.6s to improve throughput.
        - Disabled braking resistor monitoring due to nuisance trips (temporary).
        - Added alarm mapping for DriveFault_F30005.
        CHG ID: PL-A-CHG-2025W41
        """).strip()
    ),
]

def approx_tokens(s: str) -> int:
    # Crude but useful for planning: ~1 token per 4 characters
    return math.ceil(len(s) / 4)

total_tokens = sum(approx_tokens(d.text) for d in docs)
print(f"Docs loaded: {len(docs)}")
print(f"Approx prompt tokens if you paste everything: {total_tokens:,} tokens")
print("Reality check: 2 vendor PDFs (~200 pages each) can easily be 400k–1M tokens.")
print("Naive approach risk: high cost + no audit trail (no citations to exact lines).")

# Why doc-level retrieval fails for industrial review:
#
# 1. Granularity mismatch:
#    - A 200-page manual is one "document"
#    - Citing "vendor_manual.pdf" is too vague for review
#    - Reviewer needs: "page 47, section 3.2.1, lines 5-8"
#
# 2. Keyword over-matching:
#    - Doc-level matching often surfaces keyword mentions, not root causes
#    - Example: PLC code mentions "F30005" (the fault code)
#    - But the vendor manual explains the root cause
#    - Doc-level TF-IDF might rank PLC code higher (more "F30005" density)
#
# 3. No stable citations:
#    - "Top document: vendor_manual.pdf" is not actionable
#    - Human reviewer cannot verify the claim without reading the entire doc

from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity

query = "Why did Drive fault F30005 start appearing after the decel ramp change?"

vectorizer = TfidfVectorizer(stop_words="english")
X = vectorizer.fit_transform([d.text for d in docs])

qv = vectorizer.transform([query])
sims = cosine_similarity(qv, X)[0]

ranked = sorted(
    [(docs[i].source, float(sims[i])) for i in range(len(docs))],
    key=lambda x: x[1],
    reverse=True
)

print("Top docs (doc-level retrieval):")
for src, score in ranked[:3]:
    print(f"  - {src:28s} score={score:.3f}")

print("\nNaive answer (no citations):")
print("- Likely linked to a decel ramp change or braking resistor monitoring changes.")
print("- Recommended: review P1121, braking resistor config, and line voltage.")
print("\nProblem: this answer is plausible, but you cannot prove which document line drove it.")

from dataclasses import dataclass
import hashlib

@dataclass
class Chunk:
    chunk_id: str      # Hash-based stable identifier
    source: str        # Document name
    start_line: int    # First line in this chunk (1-indexed)
    end_line: int      # Last line in this chunk
    text: str          # The actual chunk text

# Line-based chunking strategy:
#
# max_lines=4: Each chunk is 4 lines maximum
# overlap=1: Chunks share 1 line with the next chunk
#
# Example document (5 lines):
#   Line 1: Fault F30005 definition
#   Line 2: Possible cause 1
#   Line 3: Possible cause 2
#   Line 4: Recommended action
#   Line 5: Document ID
#
# Chunks with max_lines=4, overlap=1:
#   Chunk 0: Lines 1-4
#   Chunk 1: Lines 4-5 (overlap at line 4)
#
# Why overlap matters:
# - Multi-line fault descriptions can span chunk boundaries
# - Overlap ensures you don't lose context at the boundary
# - Trade-off: slight storage/compute increase for better retrieval

from typing import List

def chunk_by_lines(doc: Doc, max_lines: int = 4, overlap: int = 1) -> List[Chunk]:
    """Chunk a document by lines, with overlap to preserve boundary context."""
    lines = [ln.strip() for ln in doc.text.splitlines() if ln.strip()]
    out: List[Chunk] = []
    i = 0

    while i < len(lines):
        start = i
        end = min(len(lines), i + max_lines)
        chunk_text = "\n".join(lines[start:end])

        # Generate deterministic chunk_id from content + metadata
        raw = f"{doc.source}:{start}:{end}:{chunk_text}"
        cid = hashlib.sha1(raw.encode("utf-8")).hexdigest()[:10]

        out.append(Chunk(
            chunk_id=cid,
            source=doc.source,
            start_line=start + 1,  # 1-indexed for human readability
            end_line=end,
            text=chunk_text
        ))

        if end == len(lines):
            break
        i = end - overlap  # Slide forward with overlap

    return out

# Apply chunking to all loaded documents
chunks: List[Chunk] = []
for d in docs:
    chunks.extend(chunk_by_lines(d, max_lines=4, overlap=1))

print(f"Total chunks: {len(chunks)}")
print("Sample citation key:", f"{chunks[0].source}#L{chunks[0].start_line}-L{chunks[0].end_line}|{chunks[0].chunk_id}")

# Vector embeddings and top-K retrieval:
#
# 1. Embeddings convert text into vectors (arrays of numbers)
#    - Example: "DC link overvoltage" → [0.23, -0.41, 0.87, ...]
#    - Similar text → similar vectors (measured by cosine similarity)
#
# 2. Vector stores enable fast top-K similarity search
#    - FAISS, Chroma, Pinecone, pgvector
#    - Index all chunk vectors, query with question vector
#    - Return K most similar chunks (K=4 is common)
#
# 3. TF-IDF vs OpenAI embeddings:
#    - TF-IDF: Simple, no API calls, good for demos
#    - OpenAI embeddings: Semantic, API calls, better for production
#    - Cosine similarity measures how "close" vectors are (0=orthogonal, 1=identical)

from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity

def make_citation(c: Chunk) -> str:
    return f"[{c.source}#L{c.start_line}-L{c.end_line}|{c.chunk_id}]"

chunk_vectorizer = TfidfVectorizer(stop_words="english")
CX = chunk_vectorizer.fit_transform([c.text for c in chunks])

def retrieve_chunks(query: str, k: int = 4):
    qv = chunk_vectorizer.transform([query])
    sims = cosine_similarity(qv, CX)[0]
    idxs = sims.argsort()[::-1][:k]
    return [(chunks[i], float(sims[i])) for i in idxs]

query = "Drive fault F30005 after decel ramp change braking resistor"
hits = retrieve_chunks(query, k=4)

print("Top chunks:")
for c, score in hits:
    print(f"  - {make_citation(c)} score={score:.3f}")

import json
from typing import Dict, Any

def evidence_pack(query: str, hits) -> Dict[str, Any]:
    pack = []
    for c, score in hits:
        if score <= 0:
            continue
        pack.append({
            "citation": make_citation(c),
            "score": round(score, 3),
            "text": c.text
        })
    return {"query": query, "chunks": pack}

pack = evidence_pack(query, hits)

answer_lines = [
    "Likely cause: DC link overvoltage during deceleration after the ramp was shortened.",
    "Most actionable checks:",
    "- Increase decel ramp time (e.g., P1121) or revert the recent change.",
    "- Re-enable/verify braking resistor sizing/wiring; monitoring being disabled can hide a failure.",
    "- Verify incoming line voltage is within spec.",
    "",
    "Citations:"
]
answer_lines += [f"  - {c['citation']}" for c in pack["chunks"][:3]]

answer = "\n".join(answer_lines)
print(answer)

print("\nEvidence pack (first chunk, trimmed):")
first = dict(pack["chunks"][0])
first["text"] = first["text"].splitlines()[0] + " ..."
print(json.dumps(first, indent=2))

# Minimum viable audit record:
# - timestamp: When was this query run?
# - run_id: Unique identifier for this session/shift
# - query: What did the operator ask?
# - top_k: How many chunks were retrieved?
# - citations: Which chunks were used as evidence?
# - answer_preview: What was the answer (first line for quick scan)
#
# Production additions:
# - model_version: Which LLM/embedding model was used?
# - doc_set_hash: Hash of the document set (detect version drift)
# - retrieval_scores: Similarity scores for each chunk
# - operator_id: Who ran this query?
#
# Format: JSONL (one record per line)
# - Easy to grep/process with command-line tools
# - Append-only (no risk of corrupting existing records)
# - Each line is valid JSON (can parse one line at a time)

import json
from pathlib import Path

def write_audit_log(pack, answer: str, path: str = "rag_audit_log.jsonl") -> dict:
    record = {
        "ts": "2025-12-26T00:00:00Z",
        "run_id": "demo0001",
        "query": pack["query"],
        "top_k": len(pack["chunks"]),
        "citations": [c["citation"] for c in pack["chunks"]],
        "answer_preview": answer.splitlines()[0]  # Single-line preview
    }
    Path(path).write_text(
        (Path(path).read_text() if Path(path).exists() else "") + json.dumps(record) + "\n",
        encoding="utf-8"
    )
    return record

record = write_audit_log(pack, answer)
print(json.dumps(record, indent=2))

# LlamaIndex wiring (optional): same architecture, less glue code
# This code will fall back to the TF-IDF pipeline if llama_index is not installed.

def llamaindex_available() -> bool:
    try:
        import llama_index  # noqa: F401
        return True
    except Exception:
        return False

if not llamaindex_available():
    print("llama_index not installed here → using the tested TF-IDF pipeline above.")
    print("In your environment: pip install llama-index")
else:
    # Modern LlamaIndex import style (0.10+)
    from llama_index.core import VectorStoreIndex, Document

    li_docs = [Document(text=d.text, metadata={"source": d.source}) for d in docs]
    index = VectorStoreIndex.from_documents(li_docs)

    engine = index.as_query_engine(similarity_top_k=4)
    resp = engine.query("Drive fault F30005 after decel ramp change braking resistor")

    print(resp.response)
    print("\nCitations:")
    for n in resp.source_nodes:
        src = n.node.metadata.get("source")
        print(f" - {src} score={n.score:.3f}")