Coverage for lambda / vector-ingest / handler.py: 100.00%
168 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-09-14 22:07 +0000
1"""S3-triggered ingest for the GCO vector store.
3Objects dropped under the configured corpus prefix of the
4Cluster_Shared_Bucket are chunked, embedded with the configured Bedrock
5text-embedding model, and written to the ``{project}-vector-store``
6DynamoDB global table in this (global) region; global-table replication
7fans the items out to every replica so workloads query locally.
9Wire conventions deliberately mirror the mission-memory runtime
10(``gco_mcp/mission/memory.py``): the embedding request body follows the
11Amazon Titan Text Embeddings contract (``inputText`` plus the V2-only
12``dimensions`` key, omitted for model families that reject it), and the
13vector attribute is a DynamoDB number list (``{"L": [{"N": ...}]}``)
14written through the low-level client.
16Determinism and idempotency: chunking is a pure function of the object
17bytes, and ``doc_id`` is ``sha256(key)[:16]#<chunk_index:04d>`` — so
18re-delivering an event (S3 retries at-least-once) or re-uploading an
19object overwrites the same items instead of duplicating them. Two
20consequences are documented feature limits rather than handled here:
21deleting an S3 object does not delete its items, and a shrinking object
22leaves its tail chunks behind until the corpus is re-ingested.
24Failure posture: objects are isolated — one undecodable or oversized
25object never blocks the rest of the batch — but any per-object failure
26re-raises after the batch so the async-invoke retry/DLQ machinery
27engages (succeeded objects are idempotent on the retry). The summary of
28every invocation is logged as one JSON line for operability.
30Only boto3/botocore and the standard library are used, matching every
31other GCO Lambda (no bundling step).
32"""
34from __future__ import annotations
36import hashlib
37import json
38import logging
39import os
40import urllib.parse
41from datetime import UTC, datetime
42from typing import Any
44import boto3
46# <pyflowchart-code-diagram> BEGIN - auto-inserted, do not edit
47# Generated at (UTC): 2026-09-01T14:42:56Z
48# Generated from Git commit: 89b000378ed5a912a38c06f4feab2b029936ebcc
49# Flowchart(s) generated from this file:
50# * ``lambda_handler`` -> ``diagrams/code_diagrams/lambda/vector-ingest/handler.lambda_handler.html``
51# (PNG: ``diagrams/code_diagrams/lambda/vector-ingest/handler.lambda_handler.png``)
52# Regenerate with ``SOURCE_DATE_EPOCH=<unix-seconds> GCO_DIAGRAM_SOURCE_COMMIT=<40-char-sha> python diagrams/generate.py --code-only``.
53# <pyflowchart-code-diagram> END
56logger = logging.getLogger()
57logger.setLevel(logging.INFO)
59#: Chunk-size ceiling, in characters. Paragraphs are packed greedily up
60#: to this bound; a single longer paragraph is hard-split at exactly this
61#: width. ~2000 chars keeps each chunk comfortably inside Titan's 8k-token
62#: input window while staying large enough to carry a coherent passage.
63MAX_CHUNK_CHARS = 2000
65#: Object-key suffixes routed to the plain-text paragraph chunker.
66TEXT_SUFFIXES = (".txt", ".md")
68#: Object-key suffix routed to the pre-chunked JSON-lines path.
69JSONL_SUFFIX = ".jsonl"
71#: Model-id substrings whose Titan request body accepts the V2-only
72#: ``dimensions`` key. Anything else gets the bare ``inputText`` body —
73#: V1-family models reject the key outright — and relies on the width
74#: verification below to catch a model whose default width disagrees
75#: with the deployed index.
76DIMENSIONS_CAPABLE_MODEL_MARKERS = ("titan-embed-text-v2",)
78_s3_client: Any = None
79_dynamodb_client: Any = None
80_bedrock_client: Any = None
83def _get_s3_client() -> Any:
84 global _s3_client
85 if _s3_client is None:
86 _s3_client = boto3.client("s3")
87 return _s3_client
90def _get_dynamodb_client() -> Any:
91 global _dynamodb_client
92 if _dynamodb_client is None:
93 _dynamodb_client = boto3.client("dynamodb")
94 return _dynamodb_client
97def _get_bedrock_client() -> Any:
98 global _bedrock_client
99 if _bedrock_client is None:
100 _bedrock_client = boto3.client("bedrock-runtime")
101 return _bedrock_client
104def _require_env(name: str) -> str:
105 value = os.environ.get(name, "").strip()
106 if not value:
107 raise RuntimeError(f"Required environment variable {name} is not set")
108 return value
111def _first_markdown_title(text: str) -> str | None:
112 """Return the first ATX heading's text, if the document opens with one."""
113 for line in text.splitlines():
114 stripped = line.strip()
115 if not stripped:
116 continue
117 if stripped.startswith("#"):
118 title = stripped.lstrip("#").strip()
119 return title or None
120 return None
121 return None
124def chunk_text(text: str, max_chars: int = MAX_CHUNK_CHARS) -> list[str]:
125 """Split ``text`` into deterministic ~``max_chars`` paragraph packs.
127 Paragraphs (blank-line separated) are packed greedily, joined by a
128 blank line, without ever crossing ``max_chars``; a single paragraph
129 longer than ``max_chars`` is hard-split at exactly ``max_chars``.
130 Pure function of its inputs: identical bytes always produce the
131 identical chunk list, which is what makes ``doc_id`` idempotent.
132 """
133 paragraphs = [p.strip() for p in text.replace("\r\n", "\n").split("\n\n")]
134 paragraphs = [p for p in paragraphs if p]
136 chunks: list[str] = []
137 current = ""
138 for paragraph in paragraphs:
139 pieces = (
140 [paragraph[i : i + max_chars] for i in range(0, len(paragraph), max_chars)]
141 if len(paragraph) > max_chars
142 else [paragraph]
143 )
144 for piece in pieces:
145 if not current:
146 current = piece
147 elif len(current) + 2 + len(piece) <= max_chars:
148 current = f"{current}\n\n{piece}"
149 else:
150 chunks.append(current)
151 current = piece
152 if current:
153 chunks.append(current)
154 return chunks
157def _records_from_text_object(key: str, text: str) -> list[dict[str, Any]]:
158 """Chunk a .txt/.md object into item-shaped records."""
159 title = _first_markdown_title(text)
160 records = []
161 for index, chunk in enumerate(chunk_text(text)):
162 record: dict[str, Any] = {"text": chunk, "chunk_index": index}
163 if title:
164 record["title"] = title
165 records.append(record)
166 return records
169def _records_from_jsonl_object(key: str, text: str) -> list[dict[str, Any]]:
170 """Parse a pre-chunked .jsonl object into item-shaped records.
172 Each non-empty line must be a JSON object with a non-empty string
173 ``text``; ``title`` is optional. A malformed line fails the whole
174 object (isolation stays at object granularity so a partially
175 ingested document never looks complete).
176 """
177 records: list[dict[str, Any]] = []
178 for line_number, line in enumerate(text.splitlines(), start=1):
179 if not line.strip():
180 continue
181 try:
182 payload = json.loads(line)
183 except ValueError as err:
184 raise ValueError(f"{key}: line {line_number} is not valid JSON: {err}") from err
185 if not isinstance(payload, dict):
186 raise ValueError(f"{key}: line {line_number} must be a JSON object")
187 text_value = payload.get("text")
188 if not isinstance(text_value, str) or not text_value.strip():
189 raise ValueError(f"{key}: line {line_number} needs a non-empty string 'text'")
190 if len(text_value) > MAX_CHUNK_CHARS:
191 raise ValueError(
192 f"{key}: line {line_number} text exceeds {MAX_CHUNK_CHARS} characters; "
193 "pre-chunked records must fit one chunk"
194 )
195 record: dict[str, Any] = {"text": text_value.strip(), "chunk_index": len(records)}
196 title = payload.get("title")
197 if isinstance(title, str) and title.strip():
198 record["title"] = title.strip()
199 records.append(record)
200 return records
203def _embedding_request_body(text: str, model_id: str, dimensions: int) -> str:
204 """Build the Titan-contract request body for one chunk.
206 The ``dimensions`` key is sent only to model families known to accept
207 it (Titan Text Embeddings V2); V1-style models reject unknown keys,
208 so they get the bare body and the width check below arbitrates.
209 """
210 body: dict[str, Any] = {"inputText": text}
211 if any(marker in model_id for marker in DIMENSIONS_CAPABLE_MODEL_MARKERS):
212 body["dimensions"] = dimensions
213 return json.dumps(body)
216def embed_chunk(text: str, model_id: str, dimensions: int) -> list[float]:
217 """Embed one chunk and verify the vector width.
219 A width mismatch is a hard error: the deployed index width is a
220 one-way door, and a wrong-width vector would either be rejected by
221 DynamoDB or (worse, with a misconfigured index) silently poison
222 similarity results.
223 """
224 response = _get_bedrock_client().invoke_model(
225 modelId=model_id,
226 body=_embedding_request_body(text, model_id, dimensions),
227 contentType="application/json",
228 accept="application/json",
229 )
230 payload = json.loads(response["body"].read())
231 vector = payload.get("embedding") if isinstance(payload, dict) else None
232 if not isinstance(vector, list) or not vector:
233 raise ValueError(f"embedding response from {model_id} carried no vector")
234 if len(vector) != dimensions:
235 raise ValueError(
236 f"embedding width {len(vector)} from {model_id} does not match the "
237 f"configured index width {dimensions}; check vector_store.dimensions "
238 "and vector_store.embedding_model_id (changing either after index "
239 "creation means re-creating the index and re-ingesting the corpus)"
240 )
241 return [float(value) for value in vector]
244def _number_attr(value: float) -> dict[str, str]:
245 """Render one vector component as a DynamoDB number attribute value."""
246 return {"N": repr(float(value))}
249def _put_chunk_item(
250 table_name: str,
251 *,
252 key: str,
253 record: dict[str, Any],
254 vector: list[float],
255 model_id: str,
256 content_sha256: str,
257 ingested_at: str,
258) -> str:
259 """Write one chunk item; returns its deterministic ``doc_id``."""
260 key_digest = hashlib.sha256(key.encode("utf-8")).hexdigest()[:16]
261 doc_id = f"{key_digest}#{int(record['chunk_index']):04d}"
262 item: dict[str, Any] = {
263 "doc_id": {"S": doc_id},
264 "text": {"S": record["text"]},
265 "source": {"S": key},
266 "chunk_index": {"N": str(int(record["chunk_index"]))},
267 "embedding": {"L": [_number_attr(value) for value in vector]},
268 # Provenance: vectors are only comparable to vectors from the same
269 # model, and the content hash makes re-ingest audits cheap.
270 "embedding_model_id": {"S": model_id},
271 "content_sha256": {"S": content_sha256},
272 "ingested_at": {"S": ingested_at},
273 }
274 if record.get("title"):
275 item["title"] = {"S": str(record["title"])}
276 _get_dynamodb_client().put_item(TableName=table_name, Item=item)
277 return doc_id
280def _ingest_object(bucket: str, key: str) -> dict[str, Any]:
281 """Fetch, chunk, embed, and store one S3 object."""
282 table_name = _require_env("VECTOR_STORE_TABLE_NAME")
283 model_id = _require_env("EMBEDDING_MODEL_ID")
284 dimensions = int(_require_env("EMBEDDING_DIMENSIONS"))
286 raw = _get_s3_client().get_object(Bucket=bucket, Key=key)["Body"].read()
287 content_sha256 = hashlib.sha256(raw).hexdigest()
288 text = raw.decode("utf-8")
290 lowered = key.lower()
291 if lowered.endswith(JSONL_SUFFIX):
292 records = _records_from_jsonl_object(key, text)
293 else:
294 records = _records_from_text_object(key, text)
295 if not records:
296 return {"key": key, "status": "empty", "chunks": 0}
298 ingested_at = datetime.now(UTC).isoformat()
299 doc_ids = []
300 for record in records:
301 vector = embed_chunk(record["text"], model_id, dimensions)
302 doc_ids.append(
303 _put_chunk_item(
304 table_name,
305 key=key,
306 record=record,
307 vector=vector,
308 model_id=model_id,
309 content_sha256=content_sha256,
310 ingested_at=ingested_at,
311 )
312 )
313 return {"key": key, "status": "ingested", "chunks": len(doc_ids)}
316def lambda_handler(event: dict[str, Any], context: Any) -> dict[str, Any]:
317 """Process one S3 notification event, object by object.
319 Every record is attempted (per-object isolation); a summary line is
320 logged either way; any failure re-raises afterwards so the async
321 retry/DLQ machinery sees the invocation as failed. Retries are safe:
322 ``doc_id`` is deterministic, so already-ingested objects overwrite
323 in place.
324 """
325 corpus_prefix = _require_env("CORPUS_PREFIX")
326 results: list[dict[str, Any]] = []
327 failures: list[dict[str, Any]] = []
329 for record in event.get("Records", []):
330 s3_info = record.get("s3") or {}
331 bucket = (s3_info.get("bucket") or {}).get("name")
332 raw_key = (s3_info.get("object") or {}).get("key")
333 if not bucket or not raw_key:
334 failures.append({"key": raw_key, "status": "malformed_record"})
335 continue
336 # Event keys arrive URL-encoded (spaces as '+', unicode escaped).
337 key = urllib.parse.unquote_plus(raw_key)
339 # Defense in depth: the bucket notification already filters on the
340 # prefix, but the guard keeps a mis-wired notification from
341 # ingesting arbitrary bucket contents.
342 if not key.startswith(corpus_prefix):
343 results.append({"key": key, "status": "skipped_outside_prefix", "chunks": 0})
344 continue
345 if key.endswith("/"):
346 results.append({"key": key, "status": "skipped_folder_marker", "chunks": 0})
347 continue
348 lowered = key.lower()
349 if not lowered.endswith(TEXT_SUFFIXES) and not lowered.endswith(JSONL_SUFFIX):
350 results.append({"key": key, "status": "skipped_unsupported_suffix", "chunks": 0})
351 continue
353 try:
354 results.append(_ingest_object(bucket, key))
355 except Exception as err: # noqa: BLE001 — per-object isolation, re-raised below
356 logger.exception("vector-ingest failed for s3://%s/%s", bucket, key)
357 failures.append({"key": key, "status": "failed", "error": str(err)})
359 summary = {
360 "message": "vector-ingest summary",
361 "ingested_objects": sum(1 for r in results if r["status"] == "ingested"),
362 "ingested_chunks": sum(r.get("chunks", 0) for r in results),
363 "skipped": [r for r in results if r["status"].startswith("skipped")],
364 "failures": failures,
365 }
366 logger.info(json.dumps(summary, default=str))
368 if failures:
369 raise RuntimeError(
370 f"vector-ingest failed for {len(failures)} object(s): "
371 + ", ".join(str(f["key"]) for f in failures)
372 )
373 return summary