Production RAG Infrastructure on AWS

class IndexBackend(Protocol):
    def ensure_index(
        self, *, company, name, if_exists, namespace
    ) -> str: ...

    def upsert_file(
        self, *, company, index_id, file, replace_if_newer
    ) -> None: ...

# Three implementations conform:
# - OpenSearchIndexBackend  (services/opensearch_index_backend.py)
# - AwsIndexBackend          (services/aws_index_backend.py)
# - OpenAI Vector Store path (pipelines/google_to_openai)

def ensure_retrieval(
    *, app, backend_type, company,
    drive_path, namespace, ...
):
    if backend_type == "openai":
        ensure_vs_and_docs(...)
    elif backend_type == "opensearch":
        ensure_os_index_and_docs(...)
    elif backend_type in ("bedrock", "aws"):
        ensure_aws_index_and_docs(...)

@dataclass
class OpenSearchBedrockRAG:
    """
    Intentionally generic so IRP, SCRM, CONMON, etc.
    can all reuse it by just changing `namespace`.
    """
    company: str
    namespace: str = "irp"

# opensearch.py:search_knn
if namespace:
    must_filters.append(
        {"term": {"doc.namespace": namespace}}
    )

class SourceAdapter(ABC):
    def priority(self) -> int:
        """
        1 = live API
        2 = custom/service
        3 = rag
        4 = export
        5 = manual
        """

# When the compliance engine merges results,
# a live Vanta API response (priority 1) wins over
# a RAG-retrieved document (priority 3), which wins
# over a manual CSV export (priority 5).

Pipeline         What It Does              What It Calls
─────────────────────────────────────────────────────────────────────
IRP Generator    Drafts incident           ensure_retrieval(namespace="irp")
                 response plan             → BackendExecutor.process()

SCRM Generator   Drafts supply chain       ensure_retrieval(namespace="scrm")
                 risk plan                 → BackendExecutor.process()

CONMON Generator Continuous monitoring      ensure_retrieval(namespace="conmon")
                 docs                      → BackendExecutor.process()

FedRAMP Layer 1  Requirement enrichment    Uses structured JSON input, not RAG
                                           — deliberately excluded

Custom Tests     Evidence extraction       Can query any namespace for
                 prompts                   supporting documentation

Compliance       Multi-source evidence     RAG adapter at priority=3
Engine           merge                     alongside live APIs and exports

def _chunk_text(self, text, max_tokens=500):
    words = text.split()
    current = []
    for w in words:
        current.append(w)
        if len(current) >= max_tokens:
            chunks.append(" ".join(current))
            current = []

def _split_for_embed(self, text, max_chars):
    # try splitting on a boundary
    cut = text.rfind("\n", start, end)     # prefer newline break
    if cut == -1:
        cut = text.rfind(" ", start, end)  # fallback to space break
    if cut == -1 or cut <= start + (max_chars * 0.5):
        cut = end                          # hard cut if no good boundary

# The two stages chain together in index_drive_file():
chunks = self._chunk_text(text, max_tokens=300)
max_chars = int(os.getenv("EMBED_MAX_CHARS", "12000"))
safe_chunks = []
for c in chunks:
    safe_chunks.extend(
        self._split_for_embed(c, max_chars=max_chars)
    )

Google Drive document (PDF, DOCX, Google Doc)
        |
        v
GoogleDriveIngestor.iter_files()
  -> exports Google-native docs to DOCX/XLSX/PPTX
  -> downloads binary files directly
  -> yields DriveFileRecord(content=bytes)
        |
        v
OpenSearchIndexBackend.index_drive_file()
  |
  +- freshness check: get_latest_source_modified()
  |   +- skip if unchanged (IndexResult: skipped_unchanged)
  |
  +- _extract_text(record)
  |   +- text/plain, text/markdown -> direct decode
  |   +- PDF, DOCX, PPTX, XLSX -> MarkItDown -> markdown
  |   +- DOCX fallback -> python-docx paragraph + table extraction
  |   +- binary guardrail: reject if "%PDF-" or "/FlateDecode" in output
  |
  +- _chunk_text(text, max_tokens=300)         <- Stage 1: word-window
  |
  +- _split_for_embed(chunk, max_chars=12000)  <- Stage 2: safety split
  |
  +- embed_fn(chunks)                          <- Titan Embed v2, batch=200
  |
  +- delete_docs_for_drive_file()              <- wipe stale chunks
  |
  +- os_client.index_chunks(docs)              <- bulk upsert with retry
        |
        v
OpenSearch kNN index
  mapping: doc.embedding (knn_vector, dim=1024)
  filters: doc.company (keyword), doc.namespace (keyword)
  metadata: doc.path, doc.chunk_id, doc.source_modified_at

# Score fusion algorithm
# For each document appearing in either kNN or BM25 results:

# 1. Normalize kNN scores
knn_norm = (score - min) / (max - min)      # -> [0, 1]

# 2. Normalize BM25 scores
bm25_norm = (score - min) / (max - min)     # -> [0, 1]

# 3. Fuse with configurable weights
final_score = (
    knn_norm * vector_weight
    + bm25_norm * keyword_weight
)

# 4. Sort by final_score descending, return top k

# Each hit includes scoring breakdown:
# hit["_score_detail"] = {
#     "knn_normalized": 0.8234,
#     "bm25_normalized": 0.9512,
#     "vector_weight": 0.5,
#     "keyword_weight": 0.5,
# }

# Via BackendExecutor (any pipeline):
text, meta = executor.process(
    prompt=prompt,
    instructions=instructions,
    retrieval_query=query,
    use_hybrid=True,
    vector_weight=0.5,
    keyword_weight=0.5,
)
# meta["retrieval"]["hybrid"] == True

# Directly from RAG layer:
# Pure kNN (default, unchanged behavior)
hits = rag.retrieve(
    "incident response roles and responsibilities"
)

# Hybrid (kNN + BM25)
hits = rag.retrieve(
    "KSI-CMT-01 CM-3 change advisory board",
    use_hybrid=True,
    vector_weight=0.5,
    keyword_weight=0.5,
)

retrieve(question, k=20, rerank=True, rerank_top_n=5)
                     |
  +------------------+------------------+
  | STAGE 1: Bi-Encoder Retrieval       |
  |                                     |
  | question --> Titan Embed v2 --> vec  |
  |                              |      |
  | OpenSearch <-- kNN or hybrid <+     |
  |     |                               |
  |     +-> 20 candidate hits (scored   |
  |         by cosine +/- BM25)         |
  +------------------+------------------+
                     |
  +------------------+------------------+
  | STAGE 2: LLM Reranker              |
  |                                     |
  | For each of 20 hits:                |
  |   send first 800 chars + question   |
  |   to Claude (Haiku) at temp=0       |
  |                                     |
  | Claude scores each chunk 0-10:      |
  |   10 = directly answers question    |
  |   7-9 = highly relevant detail      |
  |   4-6 = partially relevant          |
  |   1-3 = mostly irrelevant           |
  |   0 = completely irrelevant         |
  |                                     |
  | Sort by score desc -> return top 5  |
  +------------------+------------------+
                     |
  +------------------+------------------+
  | STAGE 3: Context Assembly + Answer  |
  |                                     |
  | build_context(5 hits)               |
  |   cap each hit at 4,000 chars       |
  |   cap total at 12,000 chars         |
  |                                     |
  | llm_client.chat() -> final answer   |
  +-------------------------------------+

def rerank_hits(*, question, hits, llm_client,
                top_n, preview_chars=800, model_id):

    # Step 1: Build chunk previews (lines 56-63)
    # First 800 chars of content + source path per hit
    previews = [
        {"index": i, "text": hit["content"][:preview_chars],
         "source": hit["path"]}
        for i, hit in enumerate(hits)
    ]

    # Step 2: LLM relevance scoring (lines 71-78)
    # One Bedrock call, temperature=0.0, strict JSON contract:
    # [{"index": 0, "score": 8}, {"index": 1, "score": 3}, ...]

    # Step 3: Score normalization + clamping (lines 93-98)
    # Each score clamped to [0, 10] — prevents hallucinated
    # scores of 15 or -1 from distorting the ranking

    # Step 4: Sort and prune (lines 101-113)
    # Sort by rerank score desc, original order as tiebreaker
    # Only top_n survive; each gets _rerank_score attached

    # Step 5: Graceful fallback (lines 115-118)
    # If LLM call fails or returns unparseable JSON,
    # silently return original hits truncated to top_n
    # Reranking is a quality optimization, not a gate

# Full call chain — rerank propagates through all entry points:

BackendExecutor.process(rerank=True, rerank_top_n=5)
  # -> opensearch_backend_process()  (pops rerank kwargs)
  #      -> generate_from_prompt(rerank=...)  (passes through)
  #           -> retrieve(rerank=...)  (calls rerank_hits())

answer_question(rerank=True, rerank_top_n=5)
  # -> retrieve(rerank=...)  (same path)

retrieve(rerank=True, rerank_top_n=5)
  # -> rerank_hits() directly

# Configuration via dataclass fields:
rag = OpenSearchBedrockRAG(
    os_client=os_client,
    embed_client=embed_client,
    llm_client=llm_client,          # Sonnet for answers
    company="searchstax",
    namespace="irp",
    rerank_model_id="anthropic.claude-3-5-haiku-20241022-v1:0",
    rerank_preview_chars=800,
)

@staticmethod
def _build_filters(company, namespace, metadata_filters=None):
    filters = []
    if company:
        filters.append({"term": {"doc.company": company}})
    if namespace:
        filters.append({"term": {"doc.namespace": namespace}})
    for key, value in (metadata_filters or {}).items():
        filters.append(
            {"term": {f"doc.metadata.{key}": value}}
        )
    return filters

# Both search_knn() and search_hybrid() now delegate
# to _build_filters() — no duplicated filter logic

# At indexing time — tag documents with arbitrary metadata:
backend = OpenSearchIndexBackend(
    os_client=os_client,
    embed_fn=embed_fn,
    company="searchstax",
    namespace="irp",
    extra_metadata={
        "project_id": "fedramp-2024",
        "doc_type": "policy",
        "classification": "cui",
    },
)
# Every chunk now has:
# doc.metadata.project_id = "fedramp-2024"
# doc.metadata.doc_type = "policy"
# doc.metadata.classification = "cui"

# At query time — filter to matching documents:
hits = rag.retrieve(
    "incident response roles and responsibilities",
    metadata_filters={"project_id": "fedramp-2024"},
)

# Combine with hybrid search + reranking:
hits = rag.retrieve(
    "KSI-CMT-01 change advisory board",
    use_hybrid=True,
    rerank=True,
    metadata_filters={
        "project_id": "fedramp-2024",
        "doc_type": "policy",
    },
)

# Through BackendExecutor:
text, meta = executor.process(
    prompt=prompt,
    instructions=instructions,
    metadata_filters={"project_id": "A"},
)
# meta["retrieval"]["metadata_filters"] == {"project_id": "A"}

{
  "query": {
    "bool": {
      "filter": [
        {"term": {"doc.company": "searchstax"}},
        {"term": {"doc.namespace": "irp"}},
        {"term": {"doc.metadata.project_id": "A"}}
      ],
      "must": [
        {
          "knn": {
            "doc.embedding": {
              "vector": ["..."],
              "k": 40
            }
          }
        }
      ]
    }
  }
}

Level       Mechanism                       Use Case
─────────────────────────────────────────────────────────────────
Company     doc.company (always applied)    Tenant isolation:
                                            SearchStax vs. Acme

Namespace   doc.namespace (always applied)  Document collection:
                                            IRP vs. SCRM vs. CONMON

Metadata    doc.metadata.* (opt-in via      Fine-grained:
            metadata_filters)               project, doc type,
                                            classification, version

# Critical path: embed -> search -> context -> answer
# ~200ms      ~100ms     ~0ms       ~4-6s
# (Bedrock)   (OpenSearch) (local)   (Bedrock)

# 1. Aggressive boto3 timeout (bedrock_llm.py:42-46)
Config(
    read_timeout=120,
    connect_timeout=10,    # fail fast on unreachable
    retries={"max_attempts": 3, "mode": "standard"},
)

# 2. Client singleton with connection pool reuse
# OpenSearchClient, BedrockEmbeddingsClient, BedrockLLMClient
# constructed once, injected into RAG dataclass.
# First call: ~200-500ms (credential + TLS handshake)
# Subsequent: reuse warm connection

# 3. Credential caching with TTL (ttl_cache.py)
# Google Drive SA keys + Vanta tokens cached 600s
# Without: +100ms Secrets Manager round-trip per call

# 4. Over-fetch + local score cutoff (opensearch.py:229-287)
# Fetch k * fetch_mult (5x), threshold in Python
# Avoids unreliable server-side min_score

# 5. Single-text embedding for query path
# One question -> one InvokeModel call, no batch overhead
# Batch path (EMBED_BATCH_SIZE=200) reserved for ingest

# 6. Cross-region inference profile routing
modelId = os.getenv(
    "BEDROCK_EMBEDDINGS_INFERENCE_PROFILE_ID"
) or model_id
# Routes to lowest-latency region, saves 50-100ms

Mode                     Extra Hops                         Estimated Total
────────────────────────────────────────────────────────────────────────────
Pure kNN (default)       embed + search + LLM               ~5-8s

Hybrid                   embed + 2 searches + fusion         ~5.5-8.5s
(use_hybrid=True)        + LLM

Hybrid + Rerank          embed + 2 searches + rerank LLM     ~7-11s
                         + answer LLM


Optimization Not Taken           Reason
────────────────────────────────────────────────────────────────────────────
Response streaming               Excluded from IAM policy — deliberate
(invoke_model_with_               simplicity tradeoff. Would drop perceived
response_stream)                 TTFT from ~5s to ~500ms.

OpenSearch connection            opensearch-py defaults used. Single-digit
pool tuning                      QPS on t3.small isn't connection-bound.

Formal query-path                RunMetrics tracks ingest perf (files/s,
latency instrumentation          bytes/s) but no perf_counter on retrieve().

# For each file in the Google Drive folder:

# Step 1: COMPARE — query OpenSearch for stored timestamp
latest = get_latest_source_modified(
    drive_id=record.drive_file_id
)
if latest >= record.modified_time:
    return IndexResult(status="skipped_unchanged")  # ~5ms

# Step 2: PURGE — scoped delete of all old chunks
delete_docs_for_drive_file(
    drive_id=record.drive_file_id,
    conflicts="proceed",  # concurrent searches never blocked
    refresh=True,         # cleared state immediately visible
)

# Step 3: RE-INDEX — extract, chunk, embed, bulk-index
text = _extract_text(record)
chunks = _chunk_text(text, max_tokens=300)
safe = [_split_for_embed(c, max_chars) for c in chunks]
vectors = embed_fn(safe)  # Titan Embed v2, batch=200
os_client.index_chunks(
    docs=[
        {
            "chunk_id": f"{record.drive_file_id}:{i}",
            "source_modified_at": record.modified_time,
            # ... content, embedding, metadata
        }
        for i, vec in enumerate(vectors)
    ]
)

# Typical pipeline run log:
# files=47 indexed=3 skipped_unchanged=42
# skipped_no_text=2 total_time=38.12s

Backend       Compare Strategy         Catches               Flags
──────────────────────────────────────────────────────────────────────
OpenSearch    source_modified_at       All edits (content     replace_if_newer
              timestamp comparison     + metadata changes)    skip_if_existing

S3 (AWS)      md5 + modified_time      Content edits only     replace_if_newer
              hash comparison          (misses metadata-      skip_if_existing
                                       only changes)


Both backends conform to the same IndexBackend Protocol:

class IndexBackend(Protocol):
    def upsert_file(
        self, *, company, index_id,
        file: DriveFileRecord,
        replace_if_newer: bool = True,
    ) -> None: ...

def build_context(hits, max_total=12000, max_per_hit=4000):
    """Assemble retrieved chunks into bounded, structured context."""
    context_parts = []
    total_chars = 0

    # Hits arrive in relevance order (kNN score or rerank score)
    for i, hit in enumerate(hits):
        content = hit["content"][:max_per_hit]  # cap per chunk
        if total_chars + len(content) > max_total:
            break  # budget exhausted — least important dropped

        # Structured header: boundary marker + citation anchor
        header = (
            f"[Document {i+1}] "
            f"Score: {hit['_score']:.3f} | "
            f"Chunk: {hit['chunk_id']} | "
            f"Source: {hit['path']}"
        )
        context_parts.append(f"{header}\n{content}")
        total_chars += len(content)

    return "\n\n---\n\n".join(context_parts)

# Result: 5-8 focused chunks, ~12,000 chars total
# Highest-scoring chunk at position 1 (highest attention)
# Least important material truncated by budget (correct failure mode)

SYSTEM PROMPT:
  You are a compliance documentation assistant.
  Answer ONLY using the provided context.
  Cite sources inline using [Document N] references.
  If the context does not contain the answer, say so.

---

USER PROMPT:
  Context (the ONLY allowed facts):

  [Document 1] Score: 0.923 | Chunk: abc123:3 | Source: IRP_v2.docx
  The incident response team shall acknowledge P1 incidents
  within 15 minutes of detection. The SIEM platform generates
  automated alerts...

  ---

  [Document 2] Score: 0.891 | Chunk: def456:7 | Source: IRP_v2.docx
  Escalation procedures require notification of the CISO
  within 30 minutes for incidents classified as...

  ---

  [Document 3] Score: 0.834 | Chunk: ghi789:1 | Source: roles.docx
  The Security Operations Center (SOC) maintains 24/7...

  ---

  Question: Describe the incident response acknowledgment SLA
  and escalation procedures.

Tier    Source                    Example                    Cost
──────────────────────────────────────────────────────────────────────
1       Static client profile     {{COMPANY_NAME}}           ~0ms
        (config/profile.json)     → "SearchStax"             (instant)

2       In-memory small docs      {{SIEM_TOOL}}              ~0ms
        (pre-loaded at startup)   → "Splunk Enterprise"      (cached)

3       RAG retrieval             {{INCIDENT_SLA}}           ~300ms
        (OpenSearch kNN)          → retrieved from IRP docs  (search)

4       Explicit unknown marker   {{RETENTION_PERIOD}}       ~0ms
        (unresolvable)            → "{{UNKNOWN}}"            (flagged)


Resolution order: try tier 1, then 2, then 3, then 4.
Known values are resolved BEFORE the prompt is sent to RAG —
they never enter the context window, preserving budget for
values that genuinely require document retrieval.