1#!$VENV_GTE
2"""
3================================================================================
4dsearch-answer — Full-disk semantic search with AI-ranked answer
5================================================================================
6
7PURPOSE
8 Given a natural-language question (Greek or English), find the most relevant
9 files on disk and produce a cited answer. The "CLexpand → CLrank → CLanswer"
10 pipeline: Claude expands the query, ranks files, then reads originals to answer.
11
12USAGE
13 dsearch-answer [-deep] "query" [n_results]
14 Examples:
15 dsearch-answer "find all documents about project delays"
16 dsearch-answer -deep "evaluation of Siemens proposal"
17 dsearch-answer "water bill sfaka" 50
18
19INPUTS
20 argv[1] query string (natural language)
21 argv[2] n_results (optional, default 100 files from retrieval)
22 -deep enable multi-hop retrieval (3 extra search loops)
23
24OUTPUTS (stdout)
25 Progress log per step (timestamped)
26 Final: ANSWER section with [filename] citations + SOURCES list
27
28PIPELINE (6 steps)
29 1. Query understanding → Claude parses query into JSON:
30 core_query + variants + subject_terms + HyDE
31 2. Hybrid retrieval → Cohere embeds all queries, ChromaDB finds top 100
32 (2b) Deep loop [-deep] → Read top 10 (MMR-diverse), extract new terms, repeat
33 3. Text extraction → pdftotext / antiword / zipfile / openpyxl / OCR
34 4. Parallel ranking → 10 batches × 10 files, 6 Claude workers score 0-10
35 5. Merge + dedup → top 15 by agent score (stem dedup for email threads)
36 6. Final answer → Claude reads ORIGINAL text of top 15 → cites sources
37
38EXTERNAL RESOURCES
39 ChromaDB $CHROMADB_DIR (~997K chunks, persistent)
40 Cohere embed-multilingual-v3.0 (API) (query embedding)
41 Claude claude -p via subprocess (CLI) (Steps 1, 2b, 4, 6)
42 Env keys ~/.config/dsearch/.env (COHERE_API_KEY, etc.)
43
44RELATED FILES
45 dsearch — Lightweight sibling (no agent ranking, HTML view)
46 extract_chunks.py — Builds ChromaDB from disk files (one-time indexing)
47 cohere_embed.py — Embeds chunks with Cohere (used by extract_chunks)
48 build_chroma_from_embeddings → Assembles .npy embeddings into ChromaDB collection
49 docs/index.html — Pipeline diagram (clickable, links back to here)
50
51COST PER RUN
52 Standard: 1 Claude (expand) + 10 Claude (rank) + 1 Claude (answer) + N Cohere = ~$0.20-0.40
53 Deep: +3 Claude (loop) + 3-15 extra Cohere embeds = ~$0.30-0.60
54================================================================================
55"""
56import os, sys, re, json, time, subprocess, zipfile, unicodedata
57from pathlib import Path
58from concurrent.futures import ThreadPoolExecutor, as_completed
59from json_repair import repair_json # Fixes malformed JSON from LLM output
60import regex as re_unicode # Unicode-aware regex (better Greek word boundaries)
61
62# ── Config ────────────────────────────────────────────────────────
63# Retrieval sizing
64MAX_FILES = 100 # Top N files to extract+rank after retrieval (Step 2 output)
65BATCH_SIZE = 10 # Files per agent in Step 4 (100 files ÷ 10 = 10 agents)
66
67# Text budget limits (controls Claude prompt sizes — too big = timeout / context overflow)
68MAX_FILE_CHARS = 300000 # ~120 pages per file, safe for Claude 1M context window
69MAX_AGENT_CHARS = 150000 # ~40K tokens per agent prompt — leaves headroom for claude -p
70
71# Concurrency
72AGENT_WORKERS = 6 # Parallel claude -p subprocesses in Step 4 (balance speed vs CPU load)
73
74# ChromaDB location (built by extract_chunks.py + build_chroma_from_embeddings.py)
75CHROMA_PATH = '$CHROMADB_DIR'
76COLLECTION = 'fulldisk' # ChromaDB collection name (~997K chunks from entire disk)
77
78# Deep retrieval loop (Step 2b, only with -deep flag)
79DEEP_MAX_LOOPS = 3 # Max iterations of "read → extract terms → search again"
80DEEP_PEEK_FILES = 10 # Files read per loop to extract new search terms (MMR-selected)
81
82# ── Load env ──────────────────────────────────────────────────────
83# Loads API keys (COHERE_API_KEY etc.) from ~/.config/dsearch/.env into os.environ.
84# setdefault = don't overwrite if already set in shell.
85with open('~/.config/dsearch/.env') as f:
86 for line in f:
87 if '=' in line and not line.startswith('#'):
88 k, v = line.strip().split('=', 1)
89 os.environ.setdefault(k, v)
90
91# ── Helpers ───────────────────────────────────────────────────────
92def tprint(msg):
93 """Timestamped progress print (used for inline step progress)."""
94 print(f' [{time.strftime("%H:%M:%S")}] {msg}', flush=True)
95
96def run_claude(prompt, timeout=60):
97 """Spawn 'claude -p -' subprocess, feed prompt via stdin, return stdout.
98 Why stdin instead of argv: prompts can be >100KB, exceeding POSIX ARG_MAX.
99 Returns '[timeout]' or '[error: ...]' on failure (callers handle these strings).
100 """
101 try:
102 cp = subprocess.run(
103 ['claude', '-p', '-'],
104 capture_output=True, timeout=timeout,
105 input=prompt, text=True
106 )
107 return cp.stdout.strip()
108 except subprocess.TimeoutExpired:
109 return '[timeout]'
110 except Exception as e:
111 return f'[error: {e}]'
112
113def _file_stem(filename):
114 """Normalize filename for cross-folder deduplication.
115 Returns a lowercase 'canonical' key — used to collapse copies, backups,
116 and email threads so the same content isn't ranked twice.
117
118 Handles these real-world patterns:
119 - Version suffixes: report-v2.pdf, report-BACK.pdf, report(1).pdf → "report"
120 - Email numbering: e0275_RE__Invoice.eml → "invoice"
121 - Reply prefixes: RE__Proposal.eml, Fw__Proposal.eml → "proposal"
122 """
123 name = os.path.splitext(filename)[0]
124 # Email thread dedup: strip "e0275_" export prefix + RE__/FW__ reply markers
125 name = re.sub(r'^e\d{3,5}_', '', name)
126 name = re.sub(r'^(?:RE|Re|FW|Fw)__', '', name)
127 # Backup/version suffixes: BACK, BACK2, copy, Copy2
128 name = re.sub(r'[-_]?(?:BACK|back|copy|Copy)\d*$', '', name)
129 # Parenthesized copy counter: "file(2)" → "file"
130 name = re.sub(r'\(\d+\)$', '', name)
131 # Trailing numeric: "file.2" / "file-3" → "file"
132 name = re.sub(r'[-_.]?\d+$', '', name)
133 # Trailing v-version: "file-v2", "file.v3" → "file"
134 name = re.sub(r'[-_.]?v\d+$', '', name)
135 return name.strip('.-_ ').lower()
136
137
138def extract_text(path, max_chars=MAX_FILE_CHARS):
139 """Extract plain text from any supported file format.
140 Used by Step 3 (extract) and Step 2b (deep loop peek).
141
142 Strategy: dispatch by file extension to the right tool.
143 Returns at most max_chars of text. Never raises — errors become "[error: ...]".
144
145 Supported formats:
146 .pdf → pdftotext CLI (fast, native text); falls back to .ocr.gemini.md/.txt sidecar for scans
147 .docx → unzip + strip XML tags (no python-docx dependency)
148 .odt/.ods → same trick (content.xml inside zip)
149 .doc → antiword CLI (old Word binary format)
150 .xlsx → openpyxl (first 3 sheets, 200 rows each)
151 .xls → xlrd (old Excel binary)
152 .jpg/png/etc → OCR sidecar only (no live OCR — too slow, pre-computed by bulk_ocr_gemini)
153 .html → strip tags with regex
154 * else → read as text
155 """
156 path = str(path)
157 ext = os.path.splitext(path)[1].lower()
158 try:
159 # PDF: try text layer first; if empty (scanned PDF), use Gemini OCR sidecar
160 if ext == '.pdf':
161 cp = subprocess.run(
162 ['pdftotext', '-q', '-nopgbrk', path, '-'], # -q quiet, -nopgbrk no page breaks
163 capture_output=True, timeout=15, text=True
164 )
165 text = cp.stdout
166 # Less than 50 chars = PDF has no text layer (pure scan) → use OCR cache
167 if len(text.strip()) < 50:
168 for suffix in ['.pdf.ocr.gemini.md', '.pdf.ocr.gemini.txt']:
169 ocr = path + suffix
170 if os.path.exists(ocr):
171 text = open(ocr, errors='replace').read()
172 break
173 return text[:max_chars]
174
175 # DOCX (modern Word): actually a ZIP containing XML
176 # word/document.xml holds the text + tags — strip tags with regex
177 if ext == '.docx':
178 with zipfile.ZipFile(path) as z:
179 xml = z.read('word/document.xml').decode('utf-8', errors='replace')
180 return re.sub(r'<[^>]+>', ' ', xml)[:max_chars]
181
182 # ODT / ODS (LibreOffice): same trick, content.xml is the text source
183 if ext in ('.odt', '.ods'):
184 with zipfile.ZipFile(path) as z:
185 xml = z.read('content.xml').decode('utf-8', errors='replace')
186 return re.sub(r'<[^>]+>', ' ', xml)[:max_chars]
187
188 # DOC (pre-2007 Word binary): antiword is the only reliable CLI extractor
189 if ext == '.doc':
190 cp = subprocess.run(
191 ['antiword', path],
192 capture_output=True, timeout=10, text=True
193 )
194 return cp.stdout[:max_chars]
195
196 # XLSX (modern Excel): first 3 sheets × 200 rows, join cells with " | "
197 # data_only=True evaluates formulas to their cached values
198 if ext == '.xlsx':
199 try:
200 import openpyxl
201 wb = openpyxl.load_workbook(path, read_only=True, data_only=True)
202 lines = []
203 for ws in wb.worksheets[:3]: # Cap at 3 sheets to avoid huge files
204 for row in ws.iter_rows(max_row=200, values_only=True):
205 vals = [str(c) for c in row if c is not None]
206 if vals:
207 lines.append(' | '.join(vals))
208 wb.close()
209 return '\n'.join(lines)[:max_chars]
210 except Exception:
211 return '[xlsx extraction failed]'
212
213 # XLS (pre-2007 binary Excel): xlrd only supports .xls (not .xlsx in new versions)
214 if ext == '.xls':
215 try:
216 import xlrd
217 wb = xlrd.open_workbook(path)
218 lines = []
219 for ws in wb.sheets()[:3]:
220 for r in range(min(ws.nrows, 200)):
221 vals = [str(ws.cell_value(r, c)) for c in range(ws.ncols) if ws.cell_value(r, c)]
222 if vals:
223 lines.append(' | '.join(vals))
224 return '\n'.join(lines)[:max_chars]
225 except Exception:
226 return '[xls extraction failed]'
227
228 # Images: NO live OCR (too slow for 100 files/query).
229 # Uses pre-computed Gemini OCR sidecars (from bulk_ocr_gemini-onDATA.py).
230 # Sidecar naming: image.jpg → image.jpg.ocr.gemini.txt (preferred: .md, then .txt, then pro variants)
231 if ext in ('.jpg', '.jpeg', '.png', '.tif', '.tiff', '.bmp'):
232 for suffix in ['.ocr.gemini-pro.txt', '.ocr.gemini.txt', '.ocr.gemini.md', '.ocr.txt']:
233 ocr = path + suffix
234 if os.path.exists(ocr):
235 return open(ocr, errors='replace').read()[:max_chars]
236 return '[no OCR sidecar]'
237
238 # HTML: strip <style>/<script> blocks first (content in them is useless),
239 # then strip remaining tags, then collapse whitespace.
240 if ext in ('.html', '.htm'):
241 text = open(path, errors='replace').read()
242 text = re.sub(r'<style[^>]*>.*?</style>', '', text, flags=re.DOTALL)
243 text = re.sub(r'<script[^>]*>.*?</script>', '', text, flags=re.DOTALL)
244 text = re.sub(r'<[^>]+>', ' ', text)
245 text = re.sub(r'\s+', ' ', text)
246 return text.strip()[:max_chars]
247
248 # Fallback: treat as plain text (txt, md, csv, py, json, yml, log, etc.)
249 # errors='replace' handles non-UTF8 bytes without crashing
250 return open(path, errors='replace').read()[:max_chars]
251
252 except Exception as e:
253 # Never raise — caller (Step 3) checks for "[" prefix to detect failures
254 return f'[extraction error: {e}]'
255
256
257# ── Step 1: Query understanding ───────────────────────────────────
258# INPUT: raw user query (natural language, Greek or English)
259# OUTPUT: list of search queries — [core_query, variants..., subjects..., HyDE]
260# METHOD: pre-search for context, then Claude expands into 5-8 search queries
261# WHY: one query string produces weak embeddings; expansion + HyDE catches
262# files using different vocabulary for the same concept
263def step_understand(query):
264 print(f'\n=== STEP 1: Query understanding ===', flush=True)
265 t0 = time.time()
266
267 # Pre-search: quick corpus peek to prime Claude with domain vocabulary.
268 # We embed the raw query, get top 10 filenames, feed them to Claude as context.
269 # This helps Claude guess domain-specific terms (equipment, project codes) it
270 # wouldn't know from the query alone.
271 peek_names = ''
272 try:
273 import cohere, chromadb
274 co = cohere.ClientV2(api_key=os.environ['COHERE_API_KEY'])
275 # Cohere multilingual v3 — handles Greek + English in same model
276 # input_type='search_query' tells Cohere to use query-side embedding
277 # (asymmetric retrieval: queries and documents get different embeddings)
278 resp = co.embed(texts=[query], model='embed-multilingual-v3.0',
279 input_type='search_query', embedding_types=['float'])
280 col = chromadb.PersistentClient(path=CHROMA_PATH).get_collection(COLLECTION)
281 peek = col.query(query_embeddings=resp.embeddings.float_, n_results=10)
282 # Deduplicate filenames (same file can match multiple chunks)
283 seen = set()
284 names = []
285 for meta in peek['metadatas'][0]:
286 fn = os.path.basename(meta['path'])
287 if fn not in seen:
288 seen.add(fn)
289 names.append(fn)
290 peek_names = '\n'.join(f' - {n}' for n in names[:8])
291 tprint(f'Pre-search found {len(names)} files for context')
292 except Exception as e:
293 tprint(f'Pre-search skipped: {e}')
294
295 prompt = (
296 'You are a query understanding system for a Greek/English document search engine.\n'
297 'Parse the user query and return ONLY valid JSON with these fields:\n\n'
298 '{\n'
299 ' "core_query": "the main searchable topic — strip any instructions/meta like '
300 'ψαξε, βρες, explain, summarize etc. Keep only WHAT to search for.",\n'
301 ' "variants": ["3-4 alternative phrasings of core_query. Include Greek morphological '
302 'variants (different cases/tenses/endings) and English translations."],\n'
303 ' "subject_terms": ["3-4 related TECHNICAL topics, concepts, document types, or '
304 'specific terms that would appear in files containing the answer. Think: what other '
305 'subjects, equipment, procedures, people, or technical details would documents about '
306 'this topic ALSO mention? These catch files that are relevant but use different vocabulary."],\n'
307 ' "entities": ["key proper nouns: company names, people, project names, codes"],\n'
308 ' "hypothetical_answer": "Write a SHORT (3-5 sentence) hypothetical answer to the '
309 'question. It does NOT need to be correct — it should use the VOCABULARY and TECHNICAL '
310 'TERMS that real documents containing the answer would use. Include specific domain '
311 'terms, equipment names, procedures, document types, Greek and English terms that '
312 'you GUESS would appear in the relevant files. This is used for HyDE (Hypothetical '
313 'Document Embedding) — the text will be embedded to find documents with similar vocabulary."\n'
314 '}\n\n'
315 'IMPORTANT for subject_terms:\n'
316 '- Think about what the DOCUMENTS would contain, not just the question.\n'
317 '- If the question is about evaluating a proposal, what technical areas does the '
318 'proposal cover? What specs, equipment, standards, site conditions?\n'
319 '- If the question is about a delay, what tasks were supposed to happen? What '
320 'prerequisites, approvals, technical reviews?\n'
321 '- Include both Greek and English terms.\n\n'
322 'IMPORTANT for hypothetical_answer:\n'
323 '- Use domain-specific vocabulary: equipment models, standards, Greek technical terms.\n'
324 '- Mention plausible document types (email, report, evaluation, protocol).\n'
325 '- Include BOTH the question topic AND the likely answer topics.\n'
326 '- It is OK to guess wrong — the goal is vocabulary coverage, not correctness.\n\n'
327 'Return ONLY the JSON. No explanation.\n\n'
328 + (f'CONTEXT — files already found for this query (use these to infer domain-specific terms):\n{peek_names}\n\n' if peek_names else '')
329 + f'User query: {query}'
330 )
331 raw = run_claude(prompt, timeout=45)
332
333 # Parse Claude's JSON response — two-stage parsing for robustness.
334 # LLMs sometimes wrap JSON in markdown fences, add trailing explanations,
335 # or use invalid quote chars. json_repair handles all of this.
336 parsed = {}
337 try:
338 # Stage 1: find {...} block and try strict json.loads
339 match = re.search(r'\{.*\}', raw, re.DOTALL)
340 if match:
341 parsed = json.loads(match.group())
342 except json.JSONDecodeError:
343 pass
344 if not parsed:
345 # Stage 2: json_repair fixes common LLM mistakes (smart quotes, trailing commas, etc.)
346 try:
347 fixed = repair_json(raw)
348 parsed = json.loads(fixed) if isinstance(fixed, str) else fixed
349 # Claude sometimes returns [{...}] instead of {...} — unwrap
350 if isinstance(parsed, list) and parsed:
351 parsed = parsed[0] if isinstance(parsed[0], dict) else {}
352 if not isinstance(parsed, dict):
353 parsed = {}
354 if parsed:
355 tprint(f'JSON repaired successfully')
356 except Exception:
357 # Last resort: both parse stages failed — fall back to raw query only
358 tprint(f'JSON parse failed — using raw query as fallback')
359 parsed = {}
360
361 # Extract the 5 JSON fields (with defaults in case Claude missed some)
362 core = parsed.get('core_query', query) # Cleaned query, no instruction words
363 variants = parsed.get('variants', [])[:4] # Alternative phrasings (Greek morphology + English)
364 subjects = parsed.get('subject_terms', [])[:4] # Related topics documents might mention
365 entities = parsed.get('entities', [])[:4] # Proper nouns (not used in retrieval, just logged)
366 hyde = parsed.get('hypothetical_answer', '') # HyDE: fake answer for vocabulary matching
367
368 # Build the final search query list for Step 2.
369 # Order: core first, then variants, subjects, and HyDE last.
370 # HyDE: Hypothetical Document Embedding — a fake 3-5 sentence answer. Even if
371 # wrong, it contains the vocabulary real answer-containing docs would use.
372 # Paper: "Precise Zero-Shot Dense Retrieval without Relevance Labels" (2022)
373 search_queries = [core] + variants + subjects
374 if hyde:
375 search_queries.append(hyde)
376 # Deduplicate (case-insensitive), but skip HyDE — it's long and always unique.
377 seen = set()
378 deduped = []
379 for q in search_queries[:-1] if hyde else search_queries:
380 if q.lower() not in seen:
381 seen.add(q.lower())
382 deduped.append(q)
383 if hyde:
384 deduped.append(hyde)
385 search_queries = deduped
386
387 tprint(f'Parsed in {time.time()-t0:.1f}s')
388 tprint(f'Core: {core}')
389 if variants:
390 tprint(f'Variants ({len(variants)}):')
391 for v in variants:
392 tprint(f' → {v}')
393 if subjects:
394 tprint(f'Subject terms ({len(subjects)}):')
395 for s in subjects:
396 tprint(f' ★ {s}')
397 if entities:
398 tprint(f'Entities: {", ".join(entities)}')
399 if hyde:
400 tprint(f'HyDE ({len(hyde)} chars): {hyde[:120]}...')
401 tprint(f'Total search queries: {len(search_queries)}')
402
403 return search_queries
404
405
406# ── Step 2: Hybrid retrieval ─────────────────────────────────────
407# INPUT: list of ~5-8 search queries from Step 1
408# OUTPUT: 4 dicts keyed by file path:
409# ranked (top N tuples), file_scores (all), file_meta, file_embeds
410# METHOD: batch-embed all queries with Cohere, run each against ChromaDB,
411# merge by path (keep best score), dedup by filename stem
412# WHY: Cohere multilingual v3 understands Greek + English in one embedding space.
413# ChromaDB holds ~997K chunks — pre-indexed by extract_chunks.py + cohere_embed.py
414def step_retrieve(queries, n=MAX_FILES):
415 print(f'\n=== STEP 2: Retrieve (ChromaDB x {len(queries)} queries) ===', flush=True)
416 t0 = time.time()
417 import cohere, chromadb
418
419 co = cohere.ClientV2(api_key=os.environ['COHERE_API_KEY'])
420 # Batch embed: one API call for all queries (cheaper + faster than one-at-a-time)
421 # Returns a list of float vectors (1024-dim for embed-multilingual-v3.0)
422 resp = co.embed(
423 texts=queries, model='embed-multilingual-v3.0',
424 input_type='search_query', embedding_types=['float']
425 )
426
427 col = chromadb.PersistentClient(path=CHROMA_PATH).get_collection(COLLECTION)
428
429 # Accumulators — keyed by file path. Since multiple chunks map to the same
430 # file, we keep the BEST scoring chunk per file.
431 file_scores = {} # path → best score (higher = more relevant)
432 file_meta = {} # path → metadata (folder, ext, filename, etc.)
433 file_embeds = {} # path → embedding vector of best chunk (used by MMR in Step 2b)
434
435 # Run each query's embedding against ChromaDB (50 chunks per query)
436 # include=['embeddings'] is needed so MMR can compute diversity later
437 for i, qvec in enumerate(resp.embeddings.float_):
438 results = col.query(query_embeddings=[qvec], n_results=50,
439 include=['documents', 'metadatas', 'distances', 'embeddings'])
440 # ChromaDB returns distance (lower=better); convert to similarity score (higher=better)
441 for doc, meta, dist, emb in zip(
442 results['documents'][0], results['metadatas'][0],
443 results['distances'][0], results['embeddings'][0]):
444 path = meta['path']
445 score = 1 - dist # cosine distance → cosine similarity
446 # Keep the best chunk per file across all queries
447 if path not in file_scores or score > file_scores[path]:
448 file_scores[path] = score
449 file_meta[path] = meta
450 file_embeds[path] = emb
451
452 # Stem-based dedup: collapse email threads (RE__, FW__) and backup copies
453 # so the same content isn't ranked twice by different paths.
454 # _file_stem normalizes: "RE__Invoice.eml" and "Invoice.eml" → "invoice"
455 by_stem = {} # stem → (path, score)
456 for path, score in file_scores.items():
457 stem = _file_stem(os.path.basename(path))
458 if stem not in by_stem or score > by_stem[stem][1]:
459 by_stem[stem] = (path, score)
460 pre_dedup = len(file_scores)
461 file_scores_new = {p: s for p, s in by_stem.values()}
462 # file_meta must stay in sync with file_scores after dedup
463 file_meta = {p: file_meta[p] for p in file_scores_new if p in file_meta}
464 file_scores = file_scores_new
465 if pre_dedup != len(file_scores):
466 tprint(f'Stem dedup: {pre_dedup} → {len(file_scores)} unique stems (emails + files collapsed)')
467
468 # Sort by score descending, take top N files
469 ranked = sorted(file_scores.items(), key=lambda x: -x[1])[:n]
470 tprint(f'{len(ranked)} unique files from {len(queries)} queries in {time.time()-t0:.1f}s')
471 if ranked:
472 tprint(f'Score range: {ranked[0][1]:.3f} — {ranked[-1][1]:.3f}')
473 for i, (p, s) in enumerate(ranked[:10]):
474 tprint(f' {i+1:2d}. [{s:.3f}] {os.path.basename(p)}')
475 if len(ranked) > 10:
476 tprint(f' ... +{len(ranked)-10} more files')
477 # Return: (ranked list for Step 3, full dicts for Step 2b deep loop)
478 return [(path, score, file_meta[path]) for path, score in ranked], file_scores, file_meta, file_embeds
479
480
481def _search_chromadb(queries, file_scores, file_meta):
482 """Helper for Step 2b: run more queries, MERGE into existing dicts.
483 Used when the deep loop finds new search terms. Returns count of newly-found files.
484 """
485 import cohere, chromadb
486 co = cohere.ClientV2(api_key=os.environ['COHERE_API_KEY'])
487 resp = co.embed(
488 texts=queries, model='embed-multilingual-v3.0',
489 input_type='search_query', embedding_types=['float']
490 )
491 col = chromadb.PersistentClient(path=CHROMA_PATH).get_collection(COLLECTION)
492 new_found = 0
493 for i, qvec in enumerate(resp.embeddings.float_):
494 results = col.query(query_embeddings=[qvec], n_results=50)
495 for doc, meta, dist in zip(
496 results['documents'][0], results['metadatas'][0],
497 results['distances'][0]):
498 path = meta['path']
499 score = 1 - dist
500 if path not in file_scores:
501 new_found += 1 # Track only fresh discoveries (not score updates)
502 if path not in file_scores or score > file_scores[path]:
503 file_scores[path] = score
504 file_meta[path] = meta
505 return new_found
506
507
508# ── MMR selector ──────────────────────────────────────────────────
509# MMR = Maximal Marginal Relevance (Carbonell & Goldstein, 1998)
510# Used in Step 2b deep loop to pick files that are BOTH relevant AND different
511# from each other. Without this, we'd just read the top 10 most similar files,
512# which are often near-duplicates and give no new search terms.
513#
514# Formula per candidate: MMR = λ·relevance − (1−λ)·max_similarity_to_selected
515# λ=1.0 → pure relevance (same as sorting by score)
516# λ=0.0 → pure diversity (least similar to selected)
517# λ=0.7 → default: 70% relevance weight, 30% diversity
518def _mmr_select(file_scores, file_embeds, already_read, k=10, lam=0.7):
519 """Select k files balancing relevance (score) and diversity (MMR).
520 already_read: paths to exclude (already processed in previous loops)
521 Returns: list of (path, score) tuples in selection order.
522 """
523 import numpy as np
524
525 # Candidates = files with known embeddings, not yet read
526 candidates = [(p, s) for p, s in file_scores.items() if p not in already_read and p in file_embeds]
527 if not candidates:
528 return []
529
530 # Normalize scores to [0,1] so relevance and similarity are on the same scale
531 scores = {p: s for p, s in candidates}
532 max_s = max(scores.values())
533 min_s = min(scores.values())
534 rng = max_s - min_s if max_s > min_s else 1.0 # Avoid divide-by-zero
535
536 selected = [] # [(path, score), ...] in selection order
537 selected_embeds = [] # Matching embeddings as numpy arrays (for fast cosine sim)
538
539 # Seed with already-read files: treat them as "already in the set" so MMR
540 # picks files DIFFERENT from what we've seen in previous loops.
541 for p in already_read:
542 if p in file_embeds:
543 selected_embeds.append(np.array(file_embeds[p], dtype=np.float32))
544
545 # Greedy selection: pick one file at a time that maximizes MMR score
546 for _ in range(min(k, len(candidates))):
547 best_path = None
548 best_mmr = -999
549
550 for p, s in candidates:
551 if p in [x[0] for x in selected]:
552 continue
553 norm_score = (s - min_s) / rng # Relevance component (0-1)
554 emb = np.array(file_embeds[p], dtype=np.float32)
555
556 # Diversity component: find the MAX cosine similarity to any
557 # already-selected or already-read file (we want to minimize this).
558 max_sim = 0.0
559 for se in selected_embeds:
560 cos = float(np.dot(emb, se) / (np.linalg.norm(emb) * np.linalg.norm(se) + 1e-9))
561 if cos > max_sim:
562 max_sim = cos
563
564 # MMR formula: reward relevance, penalize similarity to what we have
565 mmr = lam * norm_score - (1 - lam) * max_sim
566 if mmr > best_mmr:
567 best_mmr = mmr
568 best_path = p
569
570 if best_path is None:
571 break
572 selected.append((best_path, scores[best_path]))
573 selected_embeds.append(np.array(file_embeds[best_path], dtype=np.float32))
574
575 return selected
576
577
578# ── Step 2b: Deep retrieval loop ──────────────────────────────────
579# INPUT: query (original), file_scores/meta/embeds (from Step 2)
580# OUTPUT: expanded ranked list (Step 2 results + files discovered by new terms)
581# METHOD: loop up to 3 times:
582# 1. MMR-pick 10 diverse files from current top results
583# 2. Extract first 500 chars of each
584# 3. Ask Claude: "what NEW search terms would find more files?"
585# 4. Embed those terms, run ChromaDB again, merge results
586# WHY: Vector search misses files that use different vocabulary. After reading
587# some results we know the domain jargon, so we can find semantic siblings.
588# Stops early if no new terms or no new files found.
589def step_deep_retrieve(query, file_scores, file_meta, file_embeds):
590 """Read diverse results (MMR), extract new search terms, search again. Repeat."""
591 import cohere, chromadb
592
593 already_read = set() # Track files we've already peeked at — don't re-read
594
595 for loop_num in range(1, DEEP_MAX_LOOPS + 1):
596 print(f'\n=== STEP 2b: Deep loop {loop_num}/{DEEP_MAX_LOOPS} ===', flush=True)
597 t0 = time.time()
598
599 # Pick 10 diverse files via MMR (not just top 10 by score, which would
600 # be near-duplicates). λ=0.7 = 70% relevance weight, 30% diversity.
601 mmr_picks = _mmr_select(file_scores, file_embeds, already_read,
602 k=DEEP_PEEK_FILES, lam=0.7)
603 if not mmr_picks:
604 tprint('No more files to read — stopping')
605 break
606
607 tprint(f'MMR selected {len(mmr_picks)} diverse files:')
608 for p, s in mmr_picks:
609 tprint(f' [{s:.3f}] {os.path.basename(p)}')
610
611 # Extract small snippets (500 chars each) — just enough for Claude to
612 # recognize domain vocabulary, not enough to overflow its context.
613 snippets = []
614 for path, score in mmr_picks:
615 already_read.add(path) # Mark so MMR won't pick again next loop
616 text = extract_text(path, max_chars=1000)
617 fname = os.path.basename(path)
618 snippets.append(f'[{score:.3f}] {fname}: {text[:500]}')
619 snippets_text = '\n\n'.join(snippets)
620
621 # Ask Claude: what NEW terms should we search for?
622 prompt = (
623 f'You are helping a document search system find MORE relevant files.\n\n'
624 f'ORIGINAL QUESTION: {query}\n\n'
625 f'We already found these {len(mmr_picks)} files:\n{snippets_text}\n\n'
626 f'Based on what these files mention, what NEW search terms should we use '
627 f'to find ADDITIONAL relevant files that our current search missed?\n'
628 f'Look for: referenced documents, technical terms, people mentioned, '
629 f'project names, equipment, specific events, file names mentioned in emails.\n\n'
630 f'Return ONLY a JSON list of 3-5 new search phrases. '
631 f'Do NOT repeat terms we already searched for.\n'
632 f'Example: ["πυρασφάλεια αισθητήρες ΕΔΣ", "ΣχολιαIntracom", "site visit report"]\n'
633 )
634 raw = run_claude(prompt, timeout=45)
635
636 # Parse Claude's JSON array of new search terms
637 match = re.search(r'\[.*\]', raw, re.DOTALL)
638 new_terms = []
639 if match:
640 try:
641 new_terms = json.loads(match.group())
642 except json.JSONDecodeError:
643 pass
644
645 # Early exit: Claude couldn't think of anything new
646 if not new_terms:
647 tprint(f'No new terms found — stopping deep search')
648 break
649
650 tprint(f'New terms from loop {loop_num}:')
651 for t in new_terms:
652 tprint(f' + {t}')
653
654 # Run the new terms through ChromaDB and merge results.
655 # _search_chromadb only increments new_found for actually-new paths.
656 prev_count = len(file_scores)
657 new_found = _search_chromadb(new_terms, file_scores, file_meta)
658 tprint(f'Found {new_found} new files ({prev_count} → {len(file_scores)}) in {time.time()-t0:.1f}s')
659
660 # Early exit: search saturated (found nothing beyond what we already have)
661 if new_found == 0:
662 tprint(f'No new files discovered — stopping deep search')
663 break
664
665 # After all loops: re-rank the full set (original + deep discoveries)
666 ranked = sorted(file_scores.items(), key=lambda x: -x[1])[:MAX_FILES]
667 tprint(f'Deep search complete: {len(file_scores)} total unique files, keeping top {len(ranked)}')
668 if ranked:
669 tprint(f'Score range after deep: {ranked[0][1]:.3f} — {ranked[-1][1]:.3f}')
670 for i, (p, s) in enumerate(ranked[:5]):
671 tprint(f' {i+1}. [{s:.3f}] {os.path.basename(p)}')
672 if len(ranked) > 5:
673 tprint(f' ... +{len(ranked)-5} more')
674 return [(path, score, file_meta[path]) for path, score in ranked]
675
676
677# ── Step 3: Extract text ─────────────────────────────────────────
678# INPUT: top N (files, score, meta) tuples from Step 2
679# OUTPUT: list of dicts with extracted text ready for Step 4 agents
680# METHOD: dispatches each file through extract_text() which handles 10+ formats
681# NOTE: sequential (not parallelized) — most extractors are fast, and parallel
682# pdftotext/antiword would just contend for CPU.
683def step_extract(files):
684 print(f'\n=== STEP 3: Extract text from {len(files)} files ===', flush=True)
685 t0 = time.time()
686 extracted = []
687 errors = 0
688 ext_counts = {} # extension → count (for diagnostics)
689 for path, score, meta in files:
690 text = extract_text(path)
691 e = meta.get('ext', os.path.splitext(path)[1])
692 ext_counts[e] = ext_counts.get(e, 0) + 1
693 # Detect extraction failures: extract_text returns "[error ...]" / "[... failed]"
694 if text.startswith('[') and ('error' in text or 'failed' in text):
695 errors += 1
696 tprint(f' FAIL: {os.path.basename(path)} — {text[:60]}')
697 # Flatten to a dict the agents will consume
698 extracted.append({
699 'path': path,
700 'score': score, # Vector similarity from Step 2
701 'filename': os.path.basename(path),
702 'folder': meta.get('folder', ''),
703 'ext': meta.get('ext', ''),
704 'text': text # Up to MAX_FILE_CHARS chars
705 })
706 ok = len(extracted) - errors
707 ext_str = ', '.join(f'{v}×{k}' for k, v in sorted(ext_counts.items(), key=lambda x: -x[1]))
708 tprint(f'Extracted {ok} OK, {errors} errors in {time.time()-t0:.1f}s')
709 tprint(f'Types: {ext_str}')
710 total_chars = sum(len(f['text']) for f in extracted)
711 tprint(f'Total text: {total_chars:,} chars ({total_chars/1000:.0f}K)')
712 return extracted
713
714
715# ── Step 4: Parallel agents ──────────────────────────────────────
716# INPUT: query + list of ~100 extracted file dicts
717# OUTPUT: list of agent_results, each with parsed rankings per batch
718# METHOD: split 100 files into batches of 10 → launch 10 Claude subprocesses
719# (max 6 concurrent) → each agent RANKS its 10 files (score 0-10) with reasons
720# WHY: Ranking 100 files in one Claude call would overflow context and lose
721# accuracy. Parallel agents = 10x speedup + better attention per file.
722# COST: 10 claude -p calls × ~30s each = ~60s wall time (6-way parallel)
723def step_agents(query, extracted):
724 n_files = len(extracted)
725 n_batches = (n_files + BATCH_SIZE - 1) // BATCH_SIZE # ceil(n/BATCH_SIZE)
726 print(f'\n=== STEP 4: {n_batches} parallel agents ({BATCH_SIZE} files each, {AGENT_WORKERS} workers) ===', flush=True)
727 t0 = time.time()
728
729 # Split into contiguous batches (preserves score-based ordering within each agent)
730 batches = []
731 for i in range(0, n_files, BATCH_SIZE):
732 batches.append(extracted[i:i+BATCH_SIZE])
733
734 # Show batch assignments
735 for i, batch in enumerate(batches):
736 fnames = ', '.join(f['filename'][:30] for f in batch[:4])
737 more = f' +{len(batch)-4}' if len(batch) > 4 else ''
738 raw_chars = sum(len(f['text']) for f in batch)
739 capped = f' → capped {MAX_AGENT_CHARS/1000:.0f}K' if raw_chars > MAX_AGENT_CHARS else ''
740 tprint(f'Agent {i+1}: {len(batch)} files ({raw_chars/1000:.0f}K chars{capped}) — {fnames}{more}')
741
742 def run_agent(batch_idx, batch):
743 """Run ONE ranking agent on ONE batch of 10 files (called in ThreadPoolExecutor)."""
744 # Build the files block, enforcing MAX_AGENT_CHARS total budget.
745 # If a single file would overflow, truncate it (head+tail) so we keep
746 # both the intro and the conclusion, which usually contain the most signal.
747 file_block = []
748 total_chars = 0
749 for j, f in enumerate(batch):
750 file_text = f['text']
751 remaining = MAX_AGENT_CHARS - total_chars
752 if remaining < 1000:
753 # Budget exhausted — include filename only (let agent know it exists)
754 file_block.append(
755 f'--- FILE {j+1}: {f["filename"]} ---\n'
756 f'Path: {f["path"]}\n'
757 f'[text omitted — agent context full]\n'
758 )
759 continue
760 if len(file_text) > remaining:
761 # Head+tail truncation: 80% start + 20% end + marker in between
762 # Preserves opening context AND closing conclusions
763 keep = remaining - 200
764 head = int(keep * 0.8)
765 tail = int(keep * 0.2)
766 file_text = file_text[:head] + '\n[...truncated...]\n' + file_text[-tail:]
767 total_chars += len(file_text)
768 file_block.append(
769 f'--- FILE {j+1}: {f["filename"]} ---\n'
770 f'Path: {f["path"]}\n'
771 f'Vector score: {f["score"]:.3f}\n'
772 f'{file_text}\n'
773 )
774 files_text = '\n'.join(file_block)
775
776 prompt = (
777 f'You are evaluation agent #{batch_idx+1}. You have {len(batch)} files.\n'
778 f'Your ONLY job is to RANK these files by relevance to the question.\n'
779 f'Do NOT answer the question. Just evaluate each file.\n\n'
780 f'USER QUESTION: {query}\n\n'
781 f'FILES:\n{files_text}\n\n'
782 f'INSTRUCTIONS:\n'
783 f'For each file, score its relevance 0-10 (10=directly answers the question, '
784 f'0=completely irrelevant). Give a one-line reason.\n\n'
785 f'FORMAT your response EXACTLY as JSON:\n'
786 f'{{"rankings": [\n'
787 f' {{"file": 1, "filename": "name.ext", "score": 8, "reason": "why relevant"}},\n'
788 f' ...\n'
789 f']}}\n'
790 f'Return ONLY valid JSON. No other text.'
791 )
792
793 # Call Claude (4-minute timeout — enough for complex Greek files)
794 result = run_claude(prompt, timeout=240)
795 # Retry strategy: if timeout, try with only half the files.
796 # Common cause: one massive file pushing Claude's thinking over the limit.
797 if result == '[timeout]':
798 half = file_block[:len(file_block)//2]
799 shorter = '\n'.join(half)
800 prompt2 = prompt.replace(files_text, shorter)
801 result = run_claude(prompt2, timeout=240)
802 return batch_idx, result, batch
803
804 # Pre-allocate result slots by index so we can fill out-of-order (agents
805 # finish in unpredictable order due to variable file sizes)
806 agent_results = [None] * len(batches)
807 with ThreadPoolExecutor(max_workers=AGENT_WORKERS) as pool:
808 # Submit all batches; pool runs up to AGENT_WORKERS concurrently
809 futures = {
810 pool.submit(run_agent, i, b): i
811 for i, b in enumerate(batches)
812 }
813 for future in as_completed(futures):
814 idx = futures[future]
815 try:
816 batch_idx, result, batch = future.result()
817 agent_results[batch_idx] = {
818 'raw': result,
819 'batch': batch
820 }
821 # Parse rankings and show preview
822 parsed = _parse_rankings(result, batch)
823 agent_results[batch_idx]['parsed'] = parsed
824 top3 = [f'{r["filename"]}({r["score"]})' for r in parsed[:3]]
825 tprint(f'Agent {batch_idx+1}/{len(batches)} done — top: {", ".join(top3)}')
826 except Exception as e:
827 tprint(f'Agent {idx+1} failed: {e}')
828 agent_results[idx] = {
829 'raw': f'[error: {e}]',
830 'batch': batches[idx],
831 'parsed': []
832 }
833
834 done = sum(1 for r in agent_results if r and r.get('parsed'))
835 tprint(f'{done}/{len(batches)} agents succeeded in {time.time()-t0:.1f}s')
836 return [r for r in agent_results if r is not None]
837
838
839def _parse_rankings(raw, batch):
840 """Extract the agent's JSON rankings from its raw stdout.
841 The prompt demands strict JSON format, but LLMs sometimes add prose before/after.
842 Strategy: grab first {...} block, parse, map file indices back to original batch.
843 Returns: list of dicts sorted by score desc, or [] on parse failure.
844 """
845 match = re.search(r'\{.*\}', raw, re.DOTALL)
846 if not match:
847 return []
848 try:
849 data = json.loads(match.group())
850 rankings = data.get('rankings', [])
851 except json.JSONDecodeError:
852 return []
853
854 # Agent uses 1-based file numbering in its output (matches the FILE 1/2/3 labels)
855 results = []
856 for r in rankings:
857 file_idx = r.get('file', 0) - 1 # Convert to 0-based Python index
858 if 0 <= file_idx < len(batch): # Bounds check (agent might hallucinate indices)
859 results.append({
860 'filename': batch[file_idx]['filename'],
861 'path': batch[file_idx]['path'],
862 'text': batch[file_idx]['text'], # Needed by Step 6 for final answer
863 'score': r.get('score', 0), # Agent's 0-10 relevance score
864 'reason': r.get('reason', ''), # Agent's one-line justification
865 'vector_score': batch[file_idx]['score'] # Original Step 2 vector similarity
866 })
867 # Sort this batch's results by agent score (Step 5 will merge across batches)
868 results.sort(key=lambda x: -x['score'])
869 return results
870
871
872# ── Step 5: Merge rankings ────────────────────────────────────────
873# INPUT: list of agent_results (one per batch)
874# OUTPUT: top 15 files by agent score (with stem dedup)
875# METHOD: flatten all agent rankings → dedup by path → dedup by stem → sort → top 15
876# WHY: Different agents can independently rank the same file if retrieval
877# produces near-duplicates. Also email threads in different folders.
878TOP_FILES_FOR_ANSWER = 15 # How many files Step 6 feeds to final answer Claude
879
880def step_merge(agent_results):
881 print(f'\n=== STEP 5: Merge rankings from all agents ===', flush=True)
882 t0 = time.time()
883
884 # Flatten: collect every ranked entry across all agent batches
885 all_ranked = []
886 for ar in agent_results:
887 for r in ar.get('parsed', []):
888 all_ranked.append(r)
889
890 # Dedup pass 1: by exact path (same file ranked by >1 agent → keep best score)
891 by_path = {}
892 for r in all_ranked:
893 p = r['path']
894 if p not in by_path or r['score'] > by_path[p]['score']:
895 by_path[p] = r
896
897 # Dedup pass 2: by filename stem (see _file_stem — collapses email threads,
898 # backup copies, version variants). Uses tie-breaker: longer text wins.
899 by_stem = {}
900 for r in by_path.values():
901 stem = _file_stem(r['filename'])
902 if stem not in by_stem or r['score'] > by_stem[stem]['score']:
903 by_stem[stem] = r
904 elif r['score'] == by_stem[stem]['score']:
905 # Tie on score → prefer the one with more text (likely more info)
906 if len(r.get('text', '')) > len(by_stem[stem].get('text', '')):
907 by_stem[stem] = r
908
909 deduped_count = len(by_path) - len(by_stem)
910 # Final sort: highest agent score first (not vector score — agents have semantic understanding)
911 merged = sorted(by_stem.values(), key=lambda x: -x['score'])
912
913 tprint(f'{len(all_ranked)} rankings from agents → {len(by_path)} unique paths → {len(by_stem)} after stem dedup ({deduped_count} dupes removed)')
914
915 # Show top files
916 for i, r in enumerate(merged[:TOP_FILES_FOR_ANSWER]):
917 tprint(f' {i+1:2d}. [agent:{r["score"]:2d} vec:{r["vector_score"]:.3f}] '
918 f'{r["filename"]} — {r["reason"][:60]}')
919
920 # Count irrelevant files (score <= 2)
921 irrelevant = sum(1 for r in merged if r['score'] <= 2)
922 tprint(f'Relevant (score>2): {len(merged)-irrelevant}, irrelevant: {irrelevant}')
923
924 return merged[:TOP_FILES_FOR_ANSWER]
925
926
927# ── Step 6: Final answer from original files ──────────────────────
928# INPUT: query + top 15 files (from Step 5)
929# OUTPUT: formatted answer string with [filename] citations + SOURCES list
930# METHOD: one big Claude call with all 15 files' FULL extracted text
931# WHY: Agents in Step 4 only saw their own 10 files. Step 6 Claude sees the
932# top 15 together and can cross-reference, resolve contradictions, and
933# synthesize a coherent answer. Uses 1M context window to fit everything.
934def step_final(query, top_files):
935 n = len(top_files)
936 print(f'\n=== STEP 6: Final answer (reading {n} original files) ===', flush=True)
937 t0 = time.time()
938 total_chars = sum(len(f['text']) for f in top_files)
939 tprint(f'Feeding {total_chars:,} chars ({total_chars/1000:.0f}K) of original file text to Claude')
940
941 # Build file block: each file includes the agent's relevance score and reason,
942 # so the final Claude knows WHY each file was chosen
943 file_block = []
944 for i, f in enumerate(top_files):
945 file_block.append(
946 f'--- FILE {i+1}: {f["filename"]} (relevance: {f["score"]}/10) ---\n'
947 f'Path: {f["path"]}\n'
948 f'Why relevant: {f["reason"]}\n'
949 f'CONTENT:\n{f["text"]}\n'
950 )
951 files_text = '\n'.join(file_block)
952
953 prompt = (
954 f'You have the {n} most relevant files to answer a question.\n'
955 f'These files were selected by evaluation agents from a larger set.\n'
956 f'Read them carefully and answer the question.\n\n'
957 f'USER QUESTION: {query}\n\n'
958 f'FILES:\n{files_text}\n\n'
959 f'INSTRUCTIONS:\n'
960 f'1. Answer the question using ONLY evidence from these files.\n'
961 f'2. Quote specific text, dates, names as evidence.\n'
962 f'3. CITE filenames [filename] inline for every claim.\n'
963 f'4. If files contradict each other, note both sides with sources.\n'
964 f'5. Be comprehensive — use all relevant files, not just the top one.\n'
965 f'6. At the end, list the files you used with a one-line summary each.\n\n'
966 f'FORMAT:\n'
967 f'ANSWER:\n[comprehensive answer with inline [filename] citations]\n\n'
968 f'SOURCES:\n'
969 f'1. filename — what it contributed\n'
970 f'2. ...\n'
971 )
972
973 # 5-minute timeout — this is the most important Claude call, give it time
974 result = run_claude(prompt, timeout=300)
975 if result == '[timeout]':
976 # Retry with top 8 files only (cuts prompt in half, often finishes in time)
977 tprint(f'WARNING: Final answer timed out at 300s — retrying with shorter prompt')
978 shorter_prompt = prompt.replace(files_text, '\n'.join(file_block[:8]))
979 result = run_claude(shorter_prompt, timeout=300)
980 tprint(f'Final answer in {time.time()-t0:.1f}s')
981 return result
982
983
984# ── Main ──────────────────────────────────────────────────────────
985# Orchestrator: parses CLI args, runs the 6 steps in order, prints final answer.
986# Exit codes: 0 = success, 1 = usage error.
987def main():
988 # Parse args — accepts "-deep" anywhere, query as first positional, n_results as second
989 args = sys.argv[1:]
990 deep_mode = False
991 if '-deep' in args:
992 deep_mode = True
993 args.remove('-deep') # Remove flag so positional args still line up
994
995 if not args:
996 print('Usage: dsearch-answer [-deep] "query" [n_results]')
997 print(' -deep Multi-hop retrieval: read results, extract new terms, search again')
998 sys.exit(1)
999
1000 query = args[0]
1001 n_results = int(args[1]) if len(args) > 1 else MAX_FILES
1002
1003 n_agents = (n_results + BATCH_SIZE - 1) // BATCH_SIZE
1004 mode_str = 'DEEP' if deep_mode else 'standard'
1005 print(f'\n{"="*56}')
1006 print(f' dsearch answer: CLexpand -> CLrank -> CLanswer')
1007 print(f' Mode: {mode_str}')
1008 print(f' Query: {query[:46]}')
1009 print(f' Files: {n_results} | Agents: {n_agents} | Workers: {AGENT_WORKERS}')
1010 print(f'{"="*56}')
1011
1012 t_start = time.time()
1013
1014 # ── Pipeline execution (see module docstring for flow) ──
1015
1016 # Step 1: Claude expands raw query into 5-8 search queries + HyDE
1017 queries = step_understand(query)
1018
1019 # Step 2: Embed queries with Cohere, search ChromaDB, get top N files by path
1020 # Returns full dicts because Step 2b (deep) needs file_embeds for MMR
1021 files, file_scores, file_meta, file_embeds = step_retrieve(queries, n=n_results)
1022 if not files and not file_scores:
1023 print('\nNo files found.')
1024 sys.exit(0)
1025
1026 # Step 2b: OPTIONAL — multi-hop retrieval (only with -deep flag)
1027 # Reads diverse results, asks Claude for new terms, searches again (up to 3 loops)
1028 if deep_mode:
1029 files = step_deep_retrieve(query, file_scores, file_meta, file_embeds)
1030
1031 if not files:
1032 print('\nNo files found.')
1033 sys.exit(0)
1034
1035 # Step 3: Extract plain text from each file (format-specific dispatch)
1036 extracted = step_extract(files)
1037
1038 # Step 4: Spawn 10 parallel Claude agents — each ranks 10 files 0-10 with reasons
1039 agent_results = step_agents(query, extracted)
1040
1041 # Step 5: Merge rankings across all agents, dedup by path+stem, take top 15
1042 top_files = step_merge(agent_results)
1043 if not top_files:
1044 print('\nNo relevant files found by agents.')
1045 sys.exit(0)
1046
1047 # Step 6: Final Claude reads top 15 originals → comprehensive answer with citations
1048 answer = step_final(query, top_files)
1049
1050 # Print result
1051 total = time.time() - t_start
1052 print(f'\n{"="*56}')
1053 print(answer)
1054 print(f'\n{"="*56}')
1055 deep_str = ' (deep)' if deep_mode else ''
1056 print(f'Total: {total:.1f}s{deep_str} | {len(queries)} queries -> {len(files)} files '
1057 f'-> {len(agent_results)} agents ranked -> top {len(top_files)} -> answer')
1058
1059
1060if __name__ == '__main__':
1061 main()