Setting the file. One moment.
Ingest · Pinecone Full Text Search · pinecone-io/skills · Skills Docs
Back to the top of the page scripts/ ingest.py
Python · 279 lines · 10 KB
so the next search call comes back empty and looks like a query bug.
17
18 This script does all three correctly:
19
20 1. Bulk-upserts in batches.
21 2. Inspects every batch result; aborts loudly on any error.
22 3. Polls `documents.search` with a sentinel query until matches appear.
23
24 You provide prepared, schema-conformant JSONL + the index name. Schema
25 validation belongs upstream;
26 this script trusts the input and focuses on getting it indexed safely.
27
28 Usage:
29
30 uv run --script ingest.py \\
31 --data processed.jsonl \\
32 --index articles \\
33 --sentinel-field body
34
35 Run `--help` for the full flag list.
36 """
37
38 from __future__ import annotations
39
40 import json
41 import os
42 import time
43 from pathlib import Path
44
45 import typer
46 from pinecone import Pinecone
47
48
49 # ---------------------------------------------------------------------------
50 # Helpers — small functions, each does one thing.
51 # ---------------------------------------------------------------------------
52
53 def load_jsonl (path: Path) -> list[ dict ]:
54 """Read a JSONL file into a list of dicts. Fail loudly on parse errors."""
55 docs: list[ dict ] = []
56 for lineno, line in enumerate (path.read_text().splitlines(), start = 1 ):
57 line = line.strip()
58 if not line:
59 continue
60 try :
61 docs.append(json.loads(line))
62 except json.JSONDecodeError as e:
63 raise typer.BadParameter( f " { path } : { lineno } : invalid JSON ( { e.msg } )" )
64 if not docs:
65 raise typer.BadParameter( f " { path } : file is empty" )
66 return docs
67
68
69 def pick_sentinel_token (docs: list[ dict ], field: str ) -> str :
70 """Pick a token from `docs[*][field]` to use as the readiness-poll query.
71
72 A sentinel just needs to match *something* in the freshly-ingested data.
73 Scan from the first doc onward and return the first whitespace-split token
74 we find — first-doc-is-special datasets (cover pages, header rows, test
75 records with empty bodies) won't make us abort.
76 """
77 for doc in docs:
78 val = doc.get(field)
79 if isinstance (val, str ) and val.strip():
80 return val.strip().split()[ 0 ]
81 sample = ", " .join( sorted (docs[ 0 ].keys())) or "(none)"
82 raise typer.BadParameter(
83 f "can't auto-pick sentinel: no document has a non-empty string in { field !r} "
84 f "(scanned all { len (docs) } record(s)). Available fields in doc[0]: { sample } . "
85 f "Either fix --sentinel-field, or pass --sentinel TEXT explicitly."
86 )
87
88
89 def upsert_batches (
90 idx,
91 namespace: str ,
92 docs: list[ dict ],
93 batch_size: int ,
94 ) -> int :
95 """Bulk-upsert in batches; abort on the first failed batch.
96
97 Why we inspect the result every time:
98 `batch_upsert` returns 202 even when individual documents fail — the
99 failures are reported in `result.errors` / `result.has_errors`.
100 """
101 upserted = 0
102 for start in range ( 0 , len (docs), batch_size):
103 batch = docs[start:start + batch_size]
104 t0 = time.time()
105 result = idx.documents.batch_upsert( namespace = namespace, documents = batch)
106 elapsed = time.time() - t0
107
108 has_errors = getattr (result, "has_errors" , False ) or getattr (result, "failed_batch_count" , 0 )
109 if has_errors:
110 for err in getattr (result, "errors" , []) or []:
111 msg = getattr (err, "error_message" , None ) or str (err)
112 typer.secho( f " batch error: { msg } " , fg = typer.colors. RED , err = True )
113 raise typer.Exit( code = 1 )
114
115 upserted += len (batch)
116 typer.echo(
117 f " batch @ { start :>6} : { len (batch) :>4} docs in { elapsed :>5.2f} s"
118 f " (total: { upserted } / { len (docs) } )"
119 )
120 return upserted
121
122
123 def poll_until_searchable (
124 idx,
125 namespace: str ,
126 sentinel_field: str ,
127 sentinel_token: str ,
128 deadline_s: int ,
129 ) -> tuple[ float , int ]:
130 """Poll `documents.search` until the sentinel query returns matches.
131
132 Why this exists:
133 After `batch_upsert` returns, Pinecone is still building the inverted
134 index. A search call that arrives during that window comes back empty.
135 Without this poll, the user sees an empty `documents.search` and
136 debugs their *query*, never noticing it was an indexing race.
137
138 Returns:
139 (seconds_elapsed, number_of_probes)
140 """
141 start = time.time()
142 deadline = start + deadline_s
143 probes = 0
144 while time.time() < deadline:
145 probes += 1
146 resp = idx.documents.search(
147 namespace = namespace,
148 top_k = 1 ,
149 score_by = [{ "type" : "text" , "field" : sentinel_field, "query" : sentinel_token}],
150 include_fields = [], # required on every search; [] = lightest payload
151 )
152 if resp.matches:
153 return time.time() - start, probes
154 time.sleep( 5 )
155
156 raise typer.Exit( code = 1 )
157
158
159 def resolve_index_with_retry (pc, name: str , * , deadline_s: int = 60 ):
160 """Resolve `pc.index(name=...)`, retrying briefly during data-plane warmup."""
161 deadline = time.time() + deadline_s
162 delay = 2.0
163 last_exc = None
164 while time.time() < deadline:
165 try :
166 return pc.index( name = name)
167 except Exception as exc:
168 last_exc = exc
169 time.sleep(delay)
170 delay = min (delay * 1.5 , 8.0 )
171 raise typer.Exit(
172 f "Could not resolve index ' { name } ' within { deadline_s } s "
173 f "(last error: { type (last_exc). __name__ } : { last_exc } ). "
174 f "Check the index exists, the API key has access to it, and that the "
175 f "data-plane host has finished provisioning (control-plane `status.ready: True` "
176 f "can lag the data plane by a few seconds)."
177 )
178
179
180 # ---------------------------------------------------------------------------
181 # CLI
182 # ---------------------------------------------------------------------------
183
184 app = typer.Typer(
185 add_completion = False ,
186 help = "Ingest a JSONL file into a Pinecone FTS index, safely." ,
187 rich_markup_mode = "rich" ,
188 )
189
190
191 @app.command ()
192 def main (
193 data: Path = typer.Option(
194 ... , "--data" , "-d" ,
195 exists = True , dir_okay = False , readable = True ,
196 help = "Path to JSONL of prepared, schema-conformant documents (one per line)." ,
197 ),
198 index: str = typer.Option(
199 ... , "--index" , "-i" ,
200 help = "Pinecone index name." ,
201 ),
202 sentinel_field: str = typer.Option(
203 ... , "--sentinel-field" , "-f" ,
204 help = "An FTS-enabled field on the index. Used for the readiness-poll query. "
205 "If you don't know which to use, pick the longest free-text field on your schema." ,
206 ),
207 namespace: str = typer.Option(
208 "__default__" , "--namespace" , "-n" ,
209 help = "Index namespace." ,
210 ),
211 batch_size: int = typer.Option(
212 50 , "--batch-size" , "-b" , min = 1 , max = 200 ,
213 help = "Documents per batch_upsert call. Reduce if your dense vectors are large "
214 "(e.g. 25 for dim=3072) and you hit payload-size errors." ,
215 ),
216 poll_deadline: int = typer.Option(
217 300 , "--poll-deadline" , min = 10 , max = 3600 ,
218 help = "Seconds to wait for docs to become searchable before giving up." ,
219 ),
220 sentinel: str | None = typer.Option(
221 None , "--sentinel" , "-s" ,
222 help = "Token used for the readiness-poll query. "
223 "Default: first word of doc[0][sentinel-field]." ,
224 ),
225 ):
226 """Bulk-ingest prepared documents into a Pinecone FTS index.
227
228 [bold]Pipeline[/bold]
229
230 1. Load JSONL.
231 2. `batch_upsert` in batches; abort on any batch error.
232 3. Poll `documents.search` with a sentinel query until matches appear.
233 4. Report timings.
234
235 [bold]Required[/bold]: PINECONE_API_KEY in the environment, an existing
236 index named [bold]--index[/bold], and prepared JSONL at [bold]--data[/bold].
237 """
238 if not os.environ.get( "PINECONE_API_KEY" ):
239 raise typer.Exit( "PINECONE_API_KEY not set in environment." )
240
241 typer.echo( f "Loading { data } ..." )
242 docs = load_jsonl(data)
243 typer.echo( f "Loaded { len (docs) } document(s)." )
244
245 if sentinel is None :
246 sentinel = pick_sentinel_token(docs, sentinel_field)
247 typer.echo( f "Sentinel: { sentinel_field } = { sentinel !r} " )
248
249 pc = Pinecone( source_tag = "pinecone_skills:full_text_search_ingest" ) # reads PINECONE_API_KEY
250 idx = resolve_index_with_retry(pc, index)
251
252 typer.echo( f " \n Upserting in batches of { batch_size } ..." )
253 t_upsert_start = time.time()
254 upserted = upsert_batches(idx, namespace, docs, batch_size)
255 upsert_seconds = time.time() - t_upsert_start
256 typer.echo( f " \n Upsert complete: { upserted } doc(s) in { upsert_seconds :.1f} s." )
257
258 typer.echo( f " \n Polling for searchability (deadline { poll_deadline } s) ..." )
259 try :
260 poll_seconds, probes = poll_until_searchable(
261 idx, namespace, sentinel_field, sentinel, poll_deadline,
262 )
263 except typer.Exit:
264 typer.secho(
265 f " \n Docs not searchable within { poll_deadline } s. "
266 f "Sentinel: { sentinel_field } = { sentinel !r} . "
267 f "Possible causes: sentinel field isn't FTS-enabled on this index; "
268 f "the upserts succeeded structurally but the documents themselves were "
269 f "rejected by the inverted-index builder; the deadline is too tight." ,
270 fg = typer.colors. RED , err = True ,
271 )
272 raise
273
274 typer.echo( f "Searchable after { poll_seconds :.1f} s ( { probes } probe(s))." )
275 typer.echo( f " \n Done — total { upsert_seconds + poll_seconds :.1f} s." )
276
277
278 if __name__ == "__main__" :
279 app()