Coverage for cli / commands / vector_cmd.py: 100.00%
106 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-09-14 22:07 +0000
1"""``gco vector`` — operate the globally replicated vector store.
3Thin click veneer over :class:`cli.vector_store.VectorStoreClient`:
4``status`` (table/replica/index state), ``ingest`` (upload documents to
5the corpus prefix; the S3-triggered Lambda does the chunking and
6embedding), and ``search`` (similarity query, optionally against a
7specific replica region). Requires the opt-in vector-store add-on
8deployed with the global stack (``vector_store.enabled`` in cdk.json).
9"""
11from __future__ import annotations
13import json
14from pathlib import Path
15from typing import Any
17import click
19_VECTOR_UNAVAILABLE_HINT = (
20 "The vector store is not available. The table, index, and ingest pipeline "
21 "ship with the global stack when vector_store.enabled is true in cdk.json "
22 "(off by default): enable it and run 'gco stacks deploy gco-global', or "
23 "wait a few minutes for the vector index to finish building after the "
24 "first deployment. 'gco vector status' shows where things stand."
25)
28def _emit_json(payload: Any, *, err: bool = False) -> None:
29 """Emit ``payload`` as a single JSON line (datetime/Path safe)."""
30 from ..output import emit_structured_document
32 emit_structured_document(
33 payload,
34 output_format="json",
35 rendered=json.dumps(payload, default=str),
36 err=err,
37 )
40def _emit_error(code: str, details: dict[str, Any] | None = None) -> None:
41 """Emit a structured error envelope to stderr."""
42 payload: dict[str, Any] = {"code": code}
43 if details is not None:
44 payload["details"] = details
45 _emit_json(payload, err=True)
48def _exit_unavailable(err: Exception) -> None:
49 """Print the deployment hint and a structured envelope, then exit 1."""
50 click.echo(_VECTOR_UNAVAILABLE_HINT, err=True)
51 _emit_error("vector_store_unavailable", {"message": str(err)})
52 raise SystemExit(1)
55def _build_client(query_region: str | None = None) -> Any:
56 """Construct the store client (SSM-lazy; free until first use)."""
57 from ..vector_store import VectorStoreClient # noqa: PLC0415
59 return VectorStoreClient(query_region=query_region)
62@click.group("vector")
63def vector() -> None:
64 """Semantic search over an S3-ingested document corpus.
66 The corpus lives in the ``{project}-vector-store`` DynamoDB global
67 table (opt-in: vector_store.enabled), replicated to every deployment
68 region so workloads and searches read locally. Drop .txt/.md/.jsonl
69 files under the corpus prefix of the cluster-shared bucket — or use
70 'gco vector ingest' — and the ingest Lambda chunks, embeds, and
71 stores them.
72 """
75@vector.command("status")
76@click.option(
77 "--region",
78 default=None,
79 help="Region whose replica to describe (default: the global region).",
80)
81@click.option(
82 "--output",
83 type=click.Choice(["json", "table"]),
84 default="json",
85 show_default=True,
86)
87def vector_status_cmd(region: str | None, output: str) -> None:
88 """Show table, replica, and vector-index state."""
89 from ..vector_store import VectorStoreUnavailableError # noqa: PLC0415
91 try:
92 status = _build_client(query_region=region).status()
93 except VectorStoreUnavailableError as err:
94 _exit_unavailable(err)
95 except Exception as err: # noqa: BLE001 — CLI boundary: envelope, don't traceback
96 _emit_error("vector_status_failed", {"message": str(err)})
97 raise SystemExit(1) from None
99 if output == "table":
100 click.echo(f" table: {status['table_name']} ({status['table_status']})")
101 click.echo(f" index: {status['index_name']} ({status['index_status']})")
102 click.echo(f" region: {status['region']}")
103 click.echo(f" items: {status['item_count']}")
104 for replica in status["replicas"]:
105 click.echo(f" replica: {replica['region']} ({replica['status']})")
106 if status["index_status"] not in ("ACTIVE",):
107 click.echo(
108 " note: searches answer ValidationException until the "
109 "index is ACTIVE (several minutes after first deploy)"
110 )
111 else:
112 _emit_json(status)
115@vector.command("ingest")
116@click.argument("files", nargs=-1, type=click.Path(exists=True, dir_okay=False, path_type=Path))
117@click.option(
118 "--demo",
119 is_flag=True,
120 help="Ingest the checkout's docs/*.md as a self-contained demo corpus.",
121)
122@click.option(
123 "--wait",
124 is_flag=True,
125 help="Wait until every uploaded document is searchable (up to 5 minutes).",
126)
127@click.option(
128 "--output",
129 type=click.Choice(["json", "table"]),
130 default="json",
131 show_default=True,
132)
133def vector_ingest_cmd(files: tuple[Path, ...], demo: bool, wait: bool, output: str) -> None:
134 """Upload FILES (.txt/.md/.jsonl) to the corpus for ingestion.
136 Uploading IS ingestion: the S3 event notification invokes the ingest
137 Lambda, which chunks, embeds, and writes the items; global-table
138 replication then fans them out to every replica region. Re-uploading
139 a file overwrites its chunks in place.
140 """
141 from ..vector_store import ( # noqa: PLC0415
142 VectorStoreUnavailableError,
143 demo_corpus_paths,
144 )
146 if demo and files:
147 _emit_error("vector_ingest_invalid_args", {"message": "--demo takes no FILES"})
148 raise SystemExit(2)
149 if not demo and not files:
150 _emit_error(
151 "vector_ingest_invalid_args",
152 {"message": "give at least one file, or use --demo"},
153 )
154 raise SystemExit(2)
156 try:
157 paths = demo_corpus_paths() if demo else list(files)
158 summary = _build_client().ingest(paths, wait_timeout_seconds=300 if wait else 0)
159 except VectorStoreUnavailableError as err:
160 _exit_unavailable(err)
161 except Exception as err: # noqa: BLE001 — CLI boundary: envelope, don't traceback
162 _emit_error("vector_ingest_failed", {"message": str(err)})
163 raise SystemExit(1) from None
165 if output == "table":
166 click.echo(f" bucket: {summary['bucket']}")
167 for key in summary["uploaded"]:
168 count = (summary.get("chunks_by_source") or {}).get(key)
169 suffix = f" ({count} chunks)" if count is not None else ""
170 click.echo(f" uploaded: {key}{suffix}")
171 if summary.get("timed_out"):
172 click.echo(
173 " note: wait timed out before every document became "
174 "searchable; ingestion continues in the background "
175 "('gco vector status' to check)"
176 )
177 elif not wait:
178 click.echo(" note: ingestion continues asynchronously (--wait to block)")
179 else:
180 _emit_json(summary)
181 if summary.get("timed_out"):
182 raise SystemExit(1)
185@vector.command("search")
186@click.argument("query")
187@click.option(
188 "--top-k",
189 type=int,
190 default=5,
191 show_default=True,
192 help="Number of similar chunks to return.",
193)
194@click.option(
195 "--source",
196 default=None,
197 help="Only return chunks from this source document (full S3 key).",
198)
199@click.option(
200 "--region",
201 default=None,
202 help="Region whose replica to query (default: the global region).",
203)
204@click.option(
205 "--output",
206 type=click.Choice(["json", "table"]),
207 default="json",
208 show_default=True,
209)
210def vector_search_cmd(
211 query: str, top_k: int, source: str | None, region: str | None, output: str
212) -> None:
213 """Search the corpus for chunks similar to QUERY."""
214 from ..vector_store import VectorStoreUnavailableError # noqa: PLC0415
216 try:
217 results = _build_client(query_region=region).search(query, top_k=top_k, source=source)
218 except VectorStoreUnavailableError as err:
219 _exit_unavailable(err)
220 except Exception as err: # noqa: BLE001 — CLI boundary: envelope, don't traceback
221 _emit_error("vector_search_failed", {"message": str(err)})
222 raise SystemExit(1) from None
224 if output == "table":
225 header = f" {'SCORE':>6} {'SOURCE':<44} {'CHUNK':>5} TEXT"
226 click.echo(header)
227 click.echo(" " + "-" * (len(header) - 2))
228 for entry in results:
229 score = entry.get("score")
230 score_text = f"{score:.3f}" if isinstance(score, (int, float)) else "-"
231 source_text = (entry.get("source") or "")[:44]
232 chunk = entry.get("chunk_index", "-")
233 text = " ".join(str(entry.get("text") or "").split())[:120]
234 click.echo(f" {score_text:>6} {source_text:<44} {chunk:>5} {text}")
235 else:
236 _emit_json({"results": results})