Add on_stream parameter to run_cli and output_or_answer functions

This commit is contained in:
2026-03-19 22:54:07 +01:00
parent 85073ebbad
commit e0649ebac6

View File

@@ -301,7 +301,7 @@ def output_or_answer(final: List[Dict], args, on_stream=None):
# ----------------------------- # -----------------------------
# Main CLI (search / answer) # Main CLI (search / answer)
# ----------------------------- # -----------------------------
def run_cli(args): def run_cli(args, on_stream=None):
# Determine mode # Determine mode
hybrid_ok = all([args.shadow_index, args.shadow_store, args.content_index, args.content_store]) hybrid_ok = all([args.shadow_index, args.shadow_store, args.content_index, args.content_store])
@@ -373,7 +373,7 @@ def run_cli(args):
rec["_score"] = rec["_ann"] rec["_score"] = rec["_ann"]
final = candidates[: max(1, min(args.k, len(candidates)))] final = candidates[: max(1, min(args.k, len(candidates)))]
return output_or_answer(final, args) return output_or_answer(final, args, on_stream=on_stream)
# HYBRID path # HYBRID path
shadow_index = faiss.read_index(args.shadow_index) shadow_index = faiss.read_index(args.shadow_index)
@@ -496,7 +496,7 @@ def run_cli(args):
# per-source cap and top-k # per-source cap and top-k
ordered = apply_per_source_cap(out_candidates, args.per_source_limit) ordered = apply_per_source_cap(out_candidates, args.per_source_limit)
final = ordered[: max(1, min(args.k, len(ordered)))] final = ordered[: max(1, min(args.k, len(ordered)))]
return output_or_answer(final, args) return output_or_answer(final, args, on_stream=on_stream)
# ----------------------------- # -----------------------------
# Rerank worker mode # Rerank worker mode