engine7 7.1.23 → 7.1.25

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,980 @@
1
+ #!/usr/bin/env python3
2
+ """
3
+ Agentic Search Pipeline — 1:1 搬运 EverOS everalgo/rank 逻辑
4
+
5
+ Pipeline:
6
+ fact-MaxSim (dense + sparse) → hybrid_full (RRF) → cluster_scoped →
7
+ aagentic_retrieve (rerank → sufficiency → multi_query/refined_query → merge)
8
+
9
+ Fallback: cluster empty → hybrid_full + rerank.
10
+
11
+ Usage:
12
+ C:/Users/24045/.openclaw/everos-venv/Scripts/python.exe agentic_search.py "丢手机"
13
+ C:/Users/24045/.openclaw/everos-venv/Scripts/python.exe agentic_search.py "丢手机" --compare
14
+ """
15
+
16
+ from __future__ import annotations
17
+
18
+ import asyncio
19
+ import json
20
+ import sqlite3
21
+ import time
22
+ from dataclasses import dataclass, field
23
+ from typing import Literal
24
+
25
+ import httpx
26
+ import lancedb
27
+
28
+ # ── Config (env-driven, override via environment variables) ────────
29
+ import os
30
+
31
+ EVEROS_URL = os.getenv("EVEROS_URL", "http://127.0.0.1:8100")
32
+ DEEPINFRA_URL = os.getenv("RERANK_URL", "https://api.deepinfra.com/v1/inference/Qwen/Qwen3-Reranker-4B")
33
+ DEEPINFRA_KEY = os.getenv("RERANK_API_KEY", "")
34
+ GLM_URL = os.getenv("LLM_BASE_URL", "https://open.bigmodel.cn/api/coding/paas/v4/chat/completions")
35
+ GLM_KEY = os.getenv("LLM_API_KEY", "")
36
+ GLM_MODEL = os.getenv("LLM_MODEL", "glm-5.2")
37
+
38
+ LANCEDB_PATH = os.getenv("LANCEDB_PATH", os.path.expanduser("~/.everos/.index/lancedb"))
39
+ SQLITE_PATH = os.getenv("SQLITE_PATH", os.path.expanduser("~/.everos/.index/sqlite/system.db"))
40
+
41
+ # These are per-request, not global — but kept as defaults
42
+ DEFAULT_USER_ID = os.getenv("EVEROS_USER_ID", "xiaomei")
43
+ DEFAULT_APP_ID = os.getenv("EVEROS_APP_ID", "xiaomei")
44
+ DEFAULT_PROJECT_ID = os.getenv("EVEROS_PROJECT_ID", "default")
45
+
46
+ # ── Hyperparameters (1:1 with everos/memory/search/agentic.py) ────────
47
+ DENSE_CANDIDATES = 50
48
+ SPARSE_CANDIDATES = 50
49
+ HYBRID_RRF_K = 40 # FIX #6: episode path uses 40, not 60
50
+ CLUSTER_BASE_CANDIDATES = 100
51
+ CLUSTER_TOP_K = 10
52
+ ROUND1_TOP_N = 30
53
+ ROUND1_RERANK_TOP_N = 10
54
+ ROUND2_CAP = 40
55
+ MULTI_QUERY_COUNT = 3
56
+ REFINEMENT_STRATEGY: str = "multi_query" # FIX #1: configurable, matches default
57
+
58
+ RERANK_INSTRUCTION = (
59
+ "Determine if the passage contains specific facts, entities "
60
+ "(names, dates, locations), or details that directly answer the question."
61
+ )
62
+
63
+ # ── Prompts (1:1 with everalgo/rank/prompts/en/) ─────────────────────
64
+
65
+ # FIX #3: Full sufficiency prompt (was simplified)
66
+ SUFFICIENCY_PROMPT = """You are an expert in information retrieval evaluation. Assess whether the retrieved documents provide sufficient information to answer the user's query.
67
+
68
+ --------------------------
69
+ User Query:
70
+ {query}
71
+
72
+ Retrieved Documents:
73
+ {retrieved_docs}
74
+ --------------------------
75
+
76
+ ### Instructions:
77
+
78
+ 1. **Analyze the Query's Needs**
79
+ - **Entities**: Who/What is being asked about?
80
+ - **Attributes**: What specific details (color, time, location, quantity)?
81
+ - **Time**: Does it ask for a specific time (absolute or relative like "last week")?
82
+
83
+ 2. **Evaluate Document Evidence**
84
+ - Check **Content**: Do the documents mention the entities and attributes?
85
+ - Check **Dates**:
86
+ - Use the `Date` field of each document.
87
+ - For relative time queries (e.g., "last week", "yesterday"), verify if document dates fall within that timeframe.
88
+ - If the query asks "When did X happen?", do you have the specific date or just a vague mention?
89
+
90
+ 3. **Judgment Logic**
91
+ - **Sufficient**: You can answer the query *completely* and *precisely* using ONLY the provided documents.
92
+ - **Insufficient**:
93
+ - The specific entity is not found.
94
+ - The entity is found, but the specific attribute (e.g., "price") is missing.
95
+ - The time reference cannot be resolved (e.g., doc says "yesterday" but has no date, or doc date doesn't match query timeframe).
96
+ - Conflicting information without resolution.
97
+
98
+ ### Output Format (strict JSON):
99
+ {{
100
+ "is_sufficient": true or false,
101
+ "reasoning": "Brief explanation. If insufficient, state WHY (e.g., 'Found X but missing date', 'No mention of Y').",
102
+ "key_information_found": ["Fact 1 (Source: Doc 1)", "Fact 2 (Source: Doc 2)"],
103
+ "missing_information": ["Specific gap 1", "Specific gap 2"]
104
+ }}
105
+
106
+ Now evaluate:"""
107
+
108
+ # FIX #4: Full multi-query prompt with 4 strategies + 3 styles
109
+ MULTI_QUERY_PROMPT = """You are an expert at query reformulation for conversational memory retrieval.
110
+ Your goal is to generate 2-3 complementary queries to find the MISSING information.
111
+
112
+ --------------------------
113
+ Original Query:
114
+ {original_query}
115
+
116
+ Key Information Found:
117
+ {key_info}
118
+
119
+ Missing Information:
120
+ {missing_info}
121
+
122
+ Retrieved Documents (Context):
123
+ {retrieved_docs}
124
+ --------------------------
125
+
126
+ ### Strategy Selection (Choose based on WHY info is missing):
127
+
128
+ **[A] Pivot / Entity Association (If entity is missing)**
129
+ - If the specific entity is not found, search for related entities or broader categories.
130
+ - Example: "manager's feedback" not found → try "performance review", "work evaluation".
131
+
132
+ **[B] Temporal Calculation (If time is missing/unclear)**
133
+ - Use `Date` from Retrieved Documents to anchor relative times.
134
+ - Example: doc dated 2024-03-15 mentions "last month" → search for "February 2024".
135
+ - Search for the *event* to find its timestamp: "When did the deadline change?"
136
+
137
+ **[C] Concept Expansion (If vocabulary mismatch)**
138
+ - Synonyms: "residence" → "living", "staying at", "moved to".
139
+ - General/Specific: "Italian cuisine" ↔ "pasta", "pizza", "restaurant".
140
+
141
+ **[D] Constraint Relaxation (If too specific)**
142
+ - If "quarterly sales report from Q3" fails, try "sales report", "Q3 results".
143
+ - Remove one constraint at a time.
144
+
145
+ ### Query Style Requirements (Use DIFFERENT styles):
146
+
147
+ 1. **Keyword Combo** (2-5 words): Key entities only. High recall.
148
+ - e.g., "project deadline", "vacation plans summer"
149
+ 2. **Natural Question** (5-10 words): Rephrased question.
150
+ - e.g., "When was the meeting scheduled?", "What was discussed about the budget?"
151
+ 3. **Hypothetical Statement** (HyDE, 5-10 words): A likely sentence in the memory.
152
+ - e.g., "We decided to postpone the launch", "The client requested changes"
153
+
154
+ ### Requirements:
155
+ - Generate 2-3 queries.
156
+ - **CRITICAL**: Use the strategies above to target the *Missing Information*.
157
+ - Keep queries SHORT and SEARCHABLE.
158
+
159
+ ### Output Format (STRICT JSON):
160
+ {{
161
+ "queries": [
162
+ "Query 1",
163
+ "Query 2",
164
+ "Query 3"
165
+ ],
166
+ "reasoning": "Strategy used for each query (e.g., Q1: Pivot, Q2: Temporal)"
167
+ }}
168
+
169
+ Now generate:
170
+ """
171
+
172
+ # FIX #1: refined_query prompt (was missing entirely)
173
+ REFINED_QUERY_PROMPT = """You are an expert at query reformulation for information retrieval.
174
+
175
+ **Task**: Generate a refined query that targets the missing information in the retrieved results.
176
+
177
+ **Original Query**:
178
+ {original_query}
179
+
180
+ **Retrieved Documents** (insufficient):
181
+ {retrieved_docs}
182
+
183
+ **Missing Information**:
184
+ {missing_info}
185
+
186
+ **Instructions**:
187
+ 1. Keep the core intent of the original query unchanged.
188
+ 2. Add specific keywords or rephrase to target the missing information.
189
+ 3. Make the query more specific and focused.
190
+ 4. The refined query should be a direct question that seeks to extract the missing facts.
191
+ 5. Do NOT change the query's meaning or make it too broad.
192
+ 6. Keep it concise (1-2 sentences maximum).
193
+
194
+ Now generate the refined query (output only the refined query, no additional text):
195
+ Original Query: {original_query}
196
+ Missing Info: {missing_info}
197
+
198
+ Refined Query:
199
+ """
200
+
201
+
202
+ # ── Data Types ───────────────────────────────────────────────────────
203
+ @dataclass
204
+ class Episode:
205
+ id: str
206
+ parent_id: str # memcell_id
207
+ subject: str
208
+ summary: str
209
+ episode: str # full content
210
+ timestamp: str
211
+ score: float = 0.0
212
+
213
+
214
+ @dataclass
215
+ class ClusterInfo:
216
+ id: str
217
+ members: list[str] # memcell_ids
218
+
219
+
220
+ @dataclass
221
+ class AgenticResult:
222
+ episodes: list[Episode]
223
+ method: str
224
+ is_sufficient: bool = False
225
+ multi_queries: list[str] = field(default_factory=list)
226
+ reasoning: str = ""
227
+ timing: dict[str, float] = field(default_factory=dict)
228
+
229
+
230
+ # ── DB Layer ─────────────────────────────────────────────────────────
231
+ class DB:
232
+ def __init__(self):
233
+ self.lance = lancedb.connect(LANCEDB_PATH)
234
+ self.episode_tbl = self.lance.open_table("episode")
235
+ self.sqlite = sqlite3.connect(SQLITE_PATH)
236
+
237
+ def get_clusters(self) -> list[ClusterInfo]:
238
+ cur = self.sqlite.execute("SELECT cluster_id, member_id FROM cluster_member")
239
+ cluster_map: dict[str, list[str]] = {}
240
+ for cid, mid in cur.fetchall():
241
+ cluster_map.setdefault(cid, []).append(mid)
242
+ return [ClusterInfo(id=cid, members=mids) for cid, mids in cluster_map.items()]
243
+
244
+ def get_all_episodes(self) -> list[Episode]:
245
+ rows = (
246
+ self.episode_tbl.search()
247
+ .where(f"owner_id = '{DEFAULT_USER_ID}'")
248
+ .limit(2000)
249
+ .to_list()
250
+ )
251
+ return [
252
+ Episode(
253
+ id=r["id"],
254
+ parent_id=r.get("parent_id", ""),
255
+ subject=r.get("subject", "") or "",
256
+ summary=r.get("summary", "") or "",
257
+ episode=r.get("episode", "") or "",
258
+ timestamp=str(r.get("timestamp", "")),
259
+ )
260
+ for r in rows
261
+ ]
262
+
263
+
264
+ # ── Search Layer ─────────────────────────────────────────────────────
265
+ # Cache: episode_id → parent_id (loaded once from LanceDB)
266
+ _ep_id_to_parent: dict[str, str] | None = None
267
+
268
+
269
+ def _load_parent_id_map() -> dict[str, str]:
270
+ """Load episode_id → parent_id from LanceDB (API doesn't return parent_id)."""
271
+ global _ep_id_to_parent
272
+ if _ep_id_to_parent is not None:
273
+ return _ep_id_to_parent
274
+ lance = lancedb.connect(LANCEDB_PATH)
275
+ tbl = lance.open_table("episode")
276
+ rows = tbl.search().where(f"owner_id = '{DEFAULT_USER_ID}'").limit(2000).to_list()
277
+ _ep_id_to_parent = {r["id"]: r.get("parent_id", "") or "" for r in rows}
278
+ return _ep_id_to_parent
279
+
280
+
281
+ def hybrid_search(query: str, top_k: int = 20, user_id: str = None, app_id: str = None, project_id: str = None) -> tuple[list[Episode], float]:
282
+ """Call EverOS hybrid search API (dense + sparse RRF fusion).
283
+ Backfills parent_id from LanceDB (API doesn't return it)."""
284
+ user_id = user_id or DEFAULT_USER_ID
285
+ app_id = app_id or DEFAULT_APP_ID
286
+ project_id = project_id or DEFAULT_PROJECT_ID
287
+ pid_map = _load_parent_id_map()
288
+ t0 = time.time()
289
+ r = httpx.post(
290
+ f"{EVEROS_URL}/api/v1/memory/search",
291
+ json={
292
+ "query": query,
293
+ "user_id": user_id,
294
+ "app_id": app_id,
295
+ "project_id": project_id,
296
+ "top_k": top_k,
297
+ "method": "hybrid",
298
+ },
299
+ timeout=60,
300
+ )
301
+ dt = time.time() - t0
302
+ data = r.json().get("data", {}).get("episodes", [])
303
+ episodes = [
304
+ Episode(
305
+ id=ep.get("id", ""),
306
+ parent_id=pid_map.get(ep.get("id", ""), ""), # backfill from LanceDB
307
+ subject=ep.get("subject", "") or "",
308
+ summary=ep.get("summary", "") or "",
309
+ episode=ep.get("episode", ep.get("summary", "")) or "",
310
+ timestamp=str(ep.get("timestamp", "")),
311
+ score=ep.get("score", 0.0),
312
+ )
313
+ for ep in data
314
+ ]
315
+ return episodes, dt
316
+
317
+
318
+ RERANK_BATCH_SIZE = 100 # DeepInfra 429s above ~200 docs per call
319
+
320
+
321
+ def rerank(query: str, episodes: list[Episode]) -> tuple[list[Episode], float]:
322
+ """DeepInfra Qwen3-Reranker cross-encoder rerank (batched to avoid 429)."""
323
+ if not episodes:
324
+ return [], 0.0
325
+ documents = [ep.episode[:500] or ep.summary[:500] or ep.subject for ep in episodes]
326
+ t0 = time.time()
327
+
328
+ all_scores: list[float] = []
329
+ for i in range(0, len(documents), RERANK_BATCH_SIZE):
330
+ batch = documents[i : i + RERANK_BATCH_SIZE]
331
+ r = httpx.post(
332
+ DEEPINFRA_URL,
333
+ headers={
334
+ "Authorization": f"Bearer {DEEPINFRA_KEY}",
335
+ "Content-Type": "application/json",
336
+ },
337
+ json={"queries": [query], "documents": batch},
338
+ timeout=30,
339
+ )
340
+ if r.status_code == 429:
341
+ # Rate limited — wait and retry this batch
342
+ time.sleep(2)
343
+ r = httpx.post(
344
+ DEEPINFRA_URL,
345
+ headers={
346
+ "Authorization": f"Bearer {DEEPINFRA_KEY}",
347
+ "Content-Type": "application/json",
348
+ },
349
+ json={"queries": [query], "documents": batch},
350
+ timeout=30,
351
+ )
352
+ batch_scores = r.json().get("scores", [])
353
+ if batch_scores and isinstance(batch_scores[0], list):
354
+ batch_scores = batch_scores[0]
355
+ all_scores.extend(batch_scores)
356
+
357
+ dt = time.time() - t0
358
+
359
+ ranked = sorted(zip(all_scores, episodes), key=lambda x: -x[0])
360
+ result = []
361
+ for score, ep in ranked:
362
+ ep.score = score
363
+ result.append(ep)
364
+ return result, dt
365
+
366
+
367
+ def llm_chat(prompt: str, temperature: float = 0.0) -> str:
368
+ """Call GLM-5.2 with thinking disabled."""
369
+ r = httpx.post(
370
+ GLM_URL,
371
+ headers={
372
+ "Authorization": f"Bearer {GLM_KEY}",
373
+ "Content-Type": "application/json",
374
+ },
375
+ json={
376
+ "model": GLM_MODEL,
377
+ "messages": [{"role": "user", "content": prompt}],
378
+ "max_tokens": 2000,
379
+ "temperature": temperature,
380
+ "thinking": {"type": "disabled"},
381
+ },
382
+ timeout=30,
383
+ )
384
+ return r.json()["choices"][0]["message"]["content"]
385
+
386
+
387
+ def extract_json(text: str) -> dict:
388
+ """Extract outermost JSON object from LLM response."""
389
+ start = text.find("{")
390
+ end = text.rfind("}")
391
+ if start == -1 or end == -1 or end < start:
392
+ raise ValueError(f"No JSON in: {text[:200]}")
393
+ return json.loads(text[start : end + 1])
394
+
395
+
396
+ # ── FIX #2: RRF Fusion ───────────────────────────────────────────────
397
+ def rrf_fusion(*ranked_lists: list[Episode], k: int = 60) -> list[Episode]:
398
+ """Reciprocal Rank Fusion over N ranked lists. score = Σ 1/(k+rank_i)."""
399
+ doc_scores: dict[str, float] = {}
400
+ doc_map: dict[str, Episode] = {}
401
+
402
+ for ranked_list in ranked_lists:
403
+ for rank, ep in enumerate(ranked_list, start=1):
404
+ if not ep.id:
405
+ continue
406
+ doc_map.setdefault(ep.id, ep)
407
+ doc_scores[ep.id] = doc_scores.get(ep.id, 0.0) + 1.0 / (k + rank)
408
+
409
+ sorted_ids = sorted(doc_scores.items(), key=lambda kv: kv[1], reverse=True)
410
+ result = []
411
+ for ep_id, score in sorted_ids:
412
+ ep = doc_map[ep_id]
413
+ ep.score = score
414
+ result.append(ep)
415
+ return result
416
+
417
+
418
+ # ── Agentic Pipeline ─────────────────────────────────────────────────
419
+ def cluster_scoped_search(
420
+ query: str,
421
+ all_episodes: list[Episode],
422
+ clusters: list[ClusterInfo],
423
+ top_k: int = CLUSTER_BASE_CANDIDATES,
424
+ ) -> list[Episode]:
425
+ """
426
+ 1. Run hybrid search for wide candidate pool
427
+ 2. Map candidates to clusters (MaxSim: cluster_score = max member score)
428
+ 3. Pick top-K clusters
429
+ 4. Expand selected clusters' full memberships
430
+ """
431
+ base_results, _ = hybrid_search(query, top_k=top_k)
432
+ if not base_results:
433
+ return []
434
+
435
+ # Build episode_id → memcell_id map
436
+ ep_to_memcell = {ep.id: ep.parent_id for ep in all_episodes if ep.parent_id}
437
+
438
+ # Build memcell_id → cluster_id map
439
+ memcell_to_cluster: dict[str, str] = {}
440
+ for cluster in clusters:
441
+ for member in cluster.members:
442
+ memcell_to_cluster[member] = cluster.id
443
+
444
+ # MaxSim: cluster_score = max(member_score) over members in base_results
445
+ cluster_scores: dict[str, float] = {}
446
+ for base_ep in base_results:
447
+ memcell_id = ep_to_memcell.get(base_ep.id, base_ep.parent_id)
448
+ if not memcell_id:
449
+ continue
450
+ cluster_id = memcell_to_cluster.get(memcell_id)
451
+ if cluster_id is None:
452
+ continue
453
+ prev = cluster_scores.get(cluster_id)
454
+ if prev is None or base_ep.score > prev:
455
+ cluster_scores[cluster_id] = base_ep.score
456
+
457
+ if not cluster_scores:
458
+ return []
459
+
460
+ selected = {
461
+ cid
462
+ for cid, _ in sorted(
463
+ cluster_scores.items(), key=lambda kv: kv[1], reverse=True
464
+ )[:CLUSTER_TOP_K]
465
+ }
466
+
467
+ member_ids: set[str] = set()
468
+ for cluster in clusters:
469
+ if cluster.id in selected:
470
+ member_ids.update(cluster.members)
471
+
472
+ return [ep for ep in all_episodes if ep.parent_id in member_ids]
473
+
474
+
475
+ def check_sufficiency(
476
+ query: str, episodes: list[Episode]
477
+ ) -> dict:
478
+ """LLM sufficiency check (temperature=0.0)."""
479
+ docs_str = format_docs(episodes)
480
+ prompt = SUFFICIENCY_PROMPT.format(query=query, retrieved_docs=docs_str)
481
+ resp = llm_chat(prompt, temperature=0.0)
482
+ return extract_json(resp)
483
+
484
+
485
+ def generate_multi_queries(
486
+ query: str,
487
+ episodes: list[Episode],
488
+ missing_info: list[str],
489
+ key_info: list[str],
490
+ ) -> list[str]:
491
+ """LLM multi-query generation (temperature=0.4)."""
492
+ docs_str = format_docs(episodes)
493
+ prompt = MULTI_QUERY_PROMPT.format(
494
+ original_query=query,
495
+ key_info=", ".join(key_info) if key_info else "N/A",
496
+ missing_info=", ".join(missing_info) if missing_info else "N/A",
497
+ retrieved_docs=docs_str,
498
+ )
499
+ resp = llm_chat(prompt, temperature=0.4)
500
+ data = extract_json(resp)
501
+ return data.get("queries", [])[:MULTI_QUERY_COUNT]
502
+
503
+
504
+ # FIX #1: refined_query strategy
505
+ def generate_refined_query(
506
+ query: str,
507
+ episodes: list[Episode],
508
+ missing_info: list[str],
509
+ ) -> str:
510
+ """LLM refined-query generation (temperature=0.3)."""
511
+ docs_str = format_docs(episodes)
512
+ prompt = REFINED_QUERY_PROMPT.format(
513
+ original_query=query,
514
+ retrieved_docs=docs_str,
515
+ missing_info=", ".join(missing_info) if missing_info else "N/A",
516
+ )
517
+ resp = llm_chat(prompt, temperature=0.3).strip()
518
+ # Strip known prefixes
519
+ for prefix in ("Refined Query:", "Output:", "Answer:", "Query:"):
520
+ if resp.startswith(prefix):
521
+ resp = resp[len(prefix):].strip()
522
+ # Validate: fallback to original if garbage
523
+ if not (5 <= len(resp) <= 300) or resp.lower() == query.lower():
524
+ return query
525
+ return resp
526
+
527
+
528
+ def format_docs(episodes: list[Episode], max_chars: int = 500) -> str:
529
+ """Format episodes as doc blocks for LLM prompt."""
530
+ if not episodes:
531
+ return "No retrieval results"
532
+ lines = []
533
+ for i, ep in enumerate(episodes, 1):
534
+ body = (ep.episode or ep.summary or ep.subject)[:max_chars]
535
+ lines.append(
536
+ f"Document {i}:\n"
537
+ f" Title: {ep.subject[:100]}\n"
538
+ f" Date: {ep.timestamp[:20]}\n"
539
+ f" Content: {body}\n"
540
+ )
541
+ return "\n".join(lines)
542
+
543
+
544
+ def dedup_by_id(episodes: list[Episode]) -> list[Episode]:
545
+ seen = set()
546
+ result = []
547
+ for ep in episodes:
548
+ if ep.id not in seen:
549
+ seen.add(ep.id)
550
+ result.append(ep)
551
+ return result
552
+
553
+
554
+ # ── Main Pipeline ────────────────────────────────────────────────────
555
+ async def agentic_search(
556
+ query: str, top_k: int = 5, strategy: str = REFINEMENT_STRATEGY
557
+ ) -> AgenticResult:
558
+ """
559
+ Full agentic pipeline 1:1 with EverOS:
560
+ 1. Load clusters + all episodes
561
+ 2. Round 1: cluster_scoped search
562
+ 3. Round 1: rerank → truncate to ROUND1_RERANK_TOP_N
563
+ 4. Sufficiency check
564
+ 5. If sufficient → return reranked[:top_k]
565
+ 6. If insufficient → Round 2 (multi_query or refined_query)
566
+ 7. Merge R1 + R2 unique, cap, final rerank
567
+ Fallback: cluster empty → hybrid + rerank
568
+ """
569
+ timing = {}
570
+ t_total = time.time()
571
+
572
+ db = DB()
573
+
574
+ # Load corpus
575
+ t0 = time.time()
576
+ clusters = db.get_clusters()
577
+ all_episodes = db.get_all_episodes()
578
+ timing["load_corpus"] = time.time() - t0
579
+
580
+ # Round 1: cluster_scoped (= acluster_retrieve)
581
+ t0 = time.time()
582
+ round1 = cluster_scoped_search(query, all_episodes, clusters)
583
+
584
+ if not round1:
585
+ # Fallback: hybrid + rerank
586
+ fb_eps, fb_time = hybrid_search(query, top_k=top_k * 2)
587
+ fb_reranked, rr_time = rerank(query, fb_eps)
588
+ fb_reranked = fb_reranked[:top_k]
589
+ timing["fallback_hybrid"] = fb_time
590
+ timing["fallback_rerank"] = rr_time
591
+ timing["total"] = time.time() - t_total
592
+ return AgenticResult(
593
+ episodes=fb_reranked,
594
+ method="fallback_hybrid",
595
+ timing=timing,
596
+ )
597
+
598
+ timing["cluster_scoped"] = time.time() - t0
599
+
600
+ # Round 1 rerank
601
+ t0 = time.time()
602
+ reranked, _ = rerank(query, round1)
603
+ reranked = reranked[:ROUND1_RERANK_TOP_N]
604
+ timing["round1_rerank"] = time.time() - t0
605
+
606
+ # Sufficiency check
607
+ t0 = time.time()
608
+ try:
609
+ sufficiency = check_sufficiency(query, reranked)
610
+ except Exception:
611
+ timing["sufficiency"] = time.time() - t0
612
+ timing["total"] = time.time() - t_total
613
+ return AgenticResult(
614
+ episodes=reranked[:top_k],
615
+ method="agentic (sufficiency_error)",
616
+ timing=timing,
617
+ )
618
+ timing["sufficiency"] = time.time() - t0
619
+
620
+ if sufficiency.get("is_sufficient", True):
621
+ timing["total"] = time.time() - t_total
622
+ return AgenticResult(
623
+ episodes=reranked[:top_k],
624
+ method="agentic (sufficient_r1)",
625
+ is_sufficient=True,
626
+ reasoning=sufficiency.get("reasoning", ""),
627
+ timing=timing,
628
+ )
629
+
630
+ # ── Round 2 ─────────────────────────────────────────────────────────
631
+ missing = sufficiency.get("missing_information", [])
632
+ key_info = sufficiency.get("key_information_found", [])
633
+
634
+ if strategy == "refined_query":
635
+ # FIX #1: refined_query path (1 query, 1 retrieve)
636
+ t0 = time.time()
637
+ try:
638
+ refined = generate_refined_query(query, reranked, missing)
639
+ except Exception:
640
+ refined = query
641
+ timing["refined_query_gen"] = time.time() - t0
642
+
643
+ # Round 2: hybrid_full (NOT cluster_scoped)
644
+ t0 = time.time()
645
+ r2_eps, _ = hybrid_search(refined, top_k=ROUND1_TOP_N)
646
+ timing["round2_search"] = time.time() - t0
647
+
648
+ # Merge: reranked + r2_unique
649
+ seen_ids = {ep.id for ep in reranked}
650
+ r2_unique = [ep for ep in r2_eps if ep.id not in seen_ids]
651
+
652
+ # Cap: len(reranked) + len(r2_kept) <= ROUND2_CAP
653
+ keep = max(0, ROUND2_CAP - len(reranked))
654
+ r2_unique = r2_unique[:keep]
655
+ merged = list(reranked) + r2_unique
656
+
657
+ # Final rerank
658
+ t0 = time.time()
659
+ final, _ = rerank(query, merged)
660
+ final = final[:top_k]
661
+ timing["final_rerank"] = time.time() - t0
662
+
663
+ timing["total"] = time.time() - t_total
664
+ return AgenticResult(
665
+ episodes=final,
666
+ method="agentic (refined_query)",
667
+ is_sufficient=False,
668
+ multi_queries=[refined],
669
+ reasoning=sufficiency.get("reasoning", ""),
670
+ timing=timing,
671
+ )
672
+
673
+ # else: multi_query path
674
+ t0 = time.time()
675
+ try:
676
+ multi_queries = generate_multi_queries(query, reranked, missing, key_info)
677
+ except Exception:
678
+ multi_queries = []
679
+ timing["multi_query_gen"] = time.time() - t0
680
+
681
+ if not multi_queries:
682
+ timing["total"] = time.time() - t_total
683
+ return AgenticResult(
684
+ episodes=reranked[:top_k],
685
+ method="agentic (multi_query_failed)",
686
+ is_sufficient=False,
687
+ reasoning=sufficiency.get("reasoning", ""),
688
+ timing=timing,
689
+ )
690
+
691
+ # Round 2: parallel hybrid_full per sub-query
692
+ t0 = time.time()
693
+ round2_lists = []
694
+ for mq in multi_queries:
695
+ r2_eps, _ = hybrid_search(mq, top_k=ROUND1_TOP_N)
696
+ round2_lists.append(r2_eps)
697
+ timing["round2_search"] = time.time() - t0
698
+
699
+ # FIX #2: RRF fuse Round 2 lists (was: flat extend)
700
+ if len(round2_lists) == 1:
701
+ fused = round2_lists[0]
702
+ else:
703
+ fused = rrf_fusion(*round2_lists, k=HYBRID_RRF_K)
704
+
705
+ # Merge: reranked + fused_unique (1:1 with _merge_truncate_rerank)
706
+ seen_ids = {ep.id for ep in reranked}
707
+ r2_unique = [ep for ep in fused if ep.id not in seen_ids]
708
+
709
+ # Cap
710
+ keep = max(0, ROUND2_CAP - len(reranked))
711
+ r2_unique = r2_unique[:keep]
712
+ merged = list(reranked) + r2_unique
713
+
714
+ # Final rerank
715
+ t0 = time.time()
716
+ final, _ = rerank(query, merged)
717
+ final = final[:top_k]
718
+ timing["final_rerank"] = time.time() - t0
719
+
720
+ timing["total"] = time.time() - t_total
721
+ return AgenticResult(
722
+ episodes=final,
723
+ method="agentic (multi_round)",
724
+ is_sufficient=False,
725
+ multi_queries=multi_queries,
726
+ reasoning=sufficiency.get("reasoning", ""),
727
+ timing=timing,
728
+ )
729
+
730
+
731
+
732
+ # -- CLI ---------------------------------------------------------------
733
+ def print_result(result: AgenticResult, query: str):
734
+ C_BOLD = "\033[1m"
735
+ C_GREEN = "\033[32m"
736
+ C_YELLOW = "\033[33m"
737
+ C_CYAN = "\033[36m"
738
+ C_DIM = "\033[2m"
739
+ C_MAGENTA = "\033[35m"
740
+ C_RESET = "\033[0m"
741
+
742
+ print(f"\n{C_BOLD}{'='*65}{C_RESET}")
743
+ if "full_agentic" in result.method or result.method.startswith("agentic"):
744
+ color = C_GREEN
745
+ elif "hybrid_agentic" in result.method:
746
+ color = C_CYAN
747
+ else:
748
+ color = C_YELLOW
749
+ print(f"{color} {result.method}{C_RESET} {C_DIM}query: {query}{C_RESET}")
750
+ if result.multi_queries:
751
+ print(f" {C_MAGENTA}queries: {result.multi_queries}{C_RESET}")
752
+ if result.reasoning:
753
+ print(f" {C_DIM}reasoning: {result.reasoning[:120]}{C_RESET}")
754
+ t = result.timing
755
+ parts = [f"{k}={v:.2f}s" for k, v in t.items() if k != "total"]
756
+ print(f" {C_DIM}{' | '.join(parts)} = {t.get('total', 0):.2f}s{C_RESET}")
757
+ print(f"{C_BOLD}{'='*65}{C_RESET}\n")
758
+
759
+ for i, ep in enumerate(result.episodes):
760
+ bar = "\u2588" * int(ep.score * 20) if ep.score > 0 else ""
761
+ print(f" {C_BOLD}[{i+1}]{C_RESET} {C_GREEN}score={ep.score:.4f}{C_RESET} {bar}")
762
+ print(f" {C_BOLD}subject:{C_RESET} {ep.subject[:80]}")
763
+ print(f" {C_DIM}summary: {ep.summary[:200]}{C_RESET}")
764
+ print()
765
+
766
+
767
+ # -- Level 2: hybrid_agentic_search -----------------------------------
768
+ async def hybrid_agentic_search(query: str, top_k: int = 5, strategy: str = "multi_query") -> AgenticResult:
769
+ """
770
+ Level 2: hybrid + rerank + LLM sufficiency + optional multi_query.
771
+ No cluster expansion. ~5-15s typical.
772
+ """
773
+ timing = {}
774
+ t_total = time.time()
775
+
776
+ # Round 1: hybrid + rerank
777
+ t0 = time.time()
778
+ eps, _ = hybrid_search(query, top_k=ROUND1_TOP_N)
779
+ if not eps:
780
+ timing["total"] = time.time() - t_total
781
+ return AgenticResult(episodes=[], method="hybrid_agentic (empty)", timing=timing)
782
+ reranked, _ = rerank(query, eps)
783
+ reranked = reranked[:ROUND1_RERANK_TOP_N]
784
+ timing["hybrid_rerank"] = time.time() - t0
785
+
786
+ # Sufficiency check
787
+ t0 = time.time()
788
+ try:
789
+ sufficiency = check_sufficiency(query, reranked)
790
+ except Exception:
791
+ timing["sufficiency"] = time.time() - t0
792
+ timing["total"] = time.time() - t_total
793
+ return AgenticResult(
794
+ episodes=reranked[:top_k],
795
+ method="hybrid_agentic (sufficiency_error)",
796
+ timing=timing,
797
+ )
798
+ timing["sufficiency"] = time.time() - t0
799
+
800
+ if sufficiency.get("is_sufficient", True):
801
+ timing["total"] = time.time() - t_total
802
+ return AgenticResult(
803
+ episodes=reranked[:top_k],
804
+ method="hybrid_agentic (sufficient_r1)",
805
+ is_sufficient=True,
806
+ reasoning=sufficiency.get("reasoning", ""),
807
+ timing=timing,
808
+ )
809
+
810
+ # -- Round 2 --
811
+ missing = sufficiency.get("missing_information", [])
812
+ key_info = sufficiency.get("key_information_found", [])
813
+
814
+ if strategy == "refined_query":
815
+ t0 = time.time()
816
+ try:
817
+ refined = generate_refined_query(query, reranked, missing)
818
+ except Exception:
819
+ refined = query
820
+ timing["refined_query_gen"] = time.time() - t0
821
+
822
+ t0 = time.time()
823
+ r2_eps, _ = hybrid_search(refined, top_k=ROUND1_TOP_N)
824
+ timing["round2_search"] = time.time() - t0
825
+
826
+ seen_ids = {ep.id for ep in reranked}
827
+ r2_unique = [ep for ep in r2_eps if ep.id not in seen_ids]
828
+ keep = max(0, ROUND2_CAP - len(reranked))
829
+ r2_unique = r2_unique[:keep]
830
+ merged = list(reranked) + r2_unique
831
+
832
+ t0 = time.time()
833
+ final, _ = rerank(query, merged)
834
+ final = final[:top_k]
835
+ timing["final_rerank"] = time.time() - t0
836
+
837
+ timing["total"] = time.time() - t_total
838
+ return AgenticResult(
839
+ episodes=final,
840
+ method="hybrid_agentic (refined_query)",
841
+ is_sufficient=False,
842
+ multi_queries=[refined],
843
+ reasoning=sufficiency.get("reasoning", ""),
844
+ timing=timing,
845
+ )
846
+
847
+ # multi_query path
848
+ t0 = time.time()
849
+ try:
850
+ multi_queries = generate_multi_queries(query, reranked, missing, key_info)
851
+ except Exception:
852
+ multi_queries = []
853
+ timing["multi_query_gen"] = time.time() - t0
854
+
855
+ if not multi_queries:
856
+ timing["total"] = time.time() - t_total
857
+ return AgenticResult(
858
+ episodes=reranked[:top_k],
859
+ method="hybrid_agentic (multi_query_failed)",
860
+ is_sufficient=False,
861
+ reasoning=sufficiency.get("reasoning", ""),
862
+ timing=timing,
863
+ )
864
+
865
+ # Parallel hybrid per sub-query
866
+ t0 = time.time()
867
+ round2_lists = []
868
+ for mq in multi_queries:
869
+ r2_eps, _ = hybrid_search(mq, top_k=ROUND1_TOP_N)
870
+ round2_lists.append(r2_eps)
871
+ timing["round2_search"] = time.time() - t0
872
+
873
+ # RRF fuse
874
+ if len(round2_lists) == 1:
875
+ fused = round2_lists[0]
876
+ else:
877
+ fused = rrf_fusion(*round2_lists, k=HYBRID_RRF_K)
878
+
879
+ # Merge
880
+ seen_ids = {ep.id for ep in reranked}
881
+ r2_unique = [ep for ep in fused if ep.id not in seen_ids]
882
+ keep = max(0, ROUND2_CAP - len(reranked))
883
+ r2_unique = r2_unique[:keep]
884
+ merged = list(reranked) + r2_unique
885
+
886
+ # Final rerank
887
+ t0 = time.time()
888
+ final, _ = rerank(query, merged)
889
+ final = final[:top_k]
890
+ timing["final_rerank"] = time.time() - t0
891
+
892
+ timing["total"] = time.time() - t_total
893
+ return AgenticResult(
894
+ episodes=final,
895
+ method="hybrid_agentic (multi_round)",
896
+ is_sufficient=False,
897
+ multi_queries=multi_queries,
898
+ reasoning=sufficiency.get("reasoning", ""),
899
+ timing=timing,
900
+ )
901
+
902
+
903
+ def main():
904
+ import argparse
905
+
906
+ parser = argparse.ArgumentParser(description="Memory Search Pipeline (3 modes)")
907
+ parser.add_argument("query", help="Search query")
908
+ parser.add_argument("--top-k", type=int, default=5)
909
+ parser.add_argument("--compare", "-c", action="store_true", help="Show all 3 modes side by side")
910
+ parser.add_argument(
911
+ "--mode", "-m",
912
+ choices=["hybrid", "hybrid_agentic", "agentic"],
913
+ default="hybrid_agentic",
914
+ help="hybrid(~2s) | hybrid_agentic(~5-15s) | agentic(~40s)",
915
+ )
916
+ parser.add_argument(
917
+ "--strategy",
918
+ choices=["multi_query", "refined_query"],
919
+ default="multi_query",
920
+ )
921
+ args = parser.parse_args()
922
+
923
+ if args.mode == "hybrid":
924
+ t0 = time.time()
925
+ eps, _ = hybrid_search(args.query, top_k=10)
926
+ reranked, _ = rerank(args.query, eps)
927
+ result = AgenticResult(
928
+ episodes=reranked[:args.top_k],
929
+ method="hybrid+rerank",
930
+ timing={"total": time.time() - t0},
931
+ )
932
+ print_result(result, args.query)
933
+
934
+ elif args.mode == "hybrid_agentic":
935
+ result = asyncio.run(hybrid_agentic_search(args.query, top_k=args.top_k, strategy=args.strategy))
936
+ print_result(result, args.query)
937
+
938
+ elif args.mode == "agentic":
939
+ result = asyncio.run(agentic_search(args.query, top_k=args.top_k, strategy=args.strategy))
940
+ print_result(result, args.query)
941
+
942
+ if args.compare:
943
+ print(f"\n{'='*65}")
944
+ print(f" --- Compare: all 3 modes ---")
945
+ print(f"{'='*65}")
946
+
947
+ # Mode 1: hybrid+rerank
948
+ t0 = time.time()
949
+ eps, _ = hybrid_search(args.query, top_k=10)
950
+ reranked, _ = rerank(args.query, eps)
951
+ dt1 = time.time() - t0
952
+ print(f"\n [Mode 1: hybrid+rerank] ({dt1:.1f}s)")
953
+ for i, ep in enumerate(reranked[:args.top_k]):
954
+ print(f" [{i+1}] score={ep.score:.4f} | {ep.subject[:65]}")
955
+
956
+ # Mode 2: hybrid_agentic
957
+ t0 = time.time()
958
+ result2 = asyncio.run(hybrid_agentic_search(args.query, top_k=args.top_k))
959
+ dt2 = time.time() - t0
960
+ print(f"\n [Mode 2: hybrid_agentic] ({dt2:.1f}s) sufficient={result2.is_sufficient}")
961
+ if result2.multi_queries:
962
+ print(f" queries: {result2.multi_queries}")
963
+ for i, ep in enumerate(result2.episodes):
964
+ print(f" [{i+1}] score={ep.score:.4f} | {ep.subject[:65]}")
965
+
966
+ # Mode 3: full agentic
967
+ t0 = time.time()
968
+ result3 = asyncio.run(agentic_search(args.query, top_k=args.top_k))
969
+ dt3 = time.time() - t0
970
+ print(f"\n [Mode 3: agentic] ({dt3:.1f}s) sufficient={result3.is_sufficient}")
971
+ if result3.multi_queries:
972
+ print(f" queries: {result3.multi_queries}")
973
+ for i, ep in enumerate(result3.episodes):
974
+ print(f" [{i+1}] score={ep.score:.4f} | {ep.subject[:65]}")
975
+
976
+ print(f"\n Timing: hybrid={dt1:.1f}s | hybrid_agentic={dt2:.1f}s | agentic={dt3:.1f}s")
977
+
978
+
979
+ if __name__ == "__main__":
980
+ main()