diff --git a/evaluations/evaluations/benchmark.py b/evaluations/evaluations/benchmark.py index 9b18f962..0d150284 100644 --- a/evaluations/evaluations/benchmark.py +++ b/evaluations/evaluations/benchmark.py @@ -1,8 +1,8 @@ import asyncio import shutil -from collections.abc import Awaitable, Callable, Mapping +from collections.abc import Callable, Mapping from pathlib import Path -from typing import Any, Literal, cast +from typing import Any, Literal, NamedTuple, cast import typer from dotenv import find_dotenv, load_dotenv @@ -17,6 +17,7 @@ from evaluations.config import ConversationInput, DatasetSpec from evaluations.datasets import DATASETS from evaluations.evaluators import ( ANSWER_EQUIVALENCE_RUBRIC, + REFUSAL_ELIGIBLE_LABELS, REFUSAL_RUBRIC, ConversationEvaluator, RefusalJudge, @@ -394,6 +395,8 @@ def _attach_relevant_uris( """ if spec.retrieval_loader is None or spec.retrieval_mapper is None: return + if not any(isinstance(case.inputs, str) for case in cases): + return corpus = spec.retrieval_loader() if limit is not None: corpus = corpus.select(range(min(limit, len(corpus)))) @@ -424,7 +427,7 @@ def _resolve_capability_config( return capability_model or config.qa.model -def _live_summary(report_cases, report_failures=()) -> dict[str, float | int] | None: +def _live_summary(report_cases, report_failures) -> dict[str, float | int] | None: """Aggregate ConversationEvaluator scores across conversations. Micro rates weight every turn equally (sums across conversations); macro @@ -447,7 +450,7 @@ def _live_summary(report_cases, report_failures=()) -> dict[str, float | int] | for failure in report_failures ) turns_total = sum(_score(case, "turns_total") for case in scored) - turns_judged = sum(_score(case, "turns_judged") or 0 for case in scored) + turns_judged = sum(_score(case, "turns_judged") for case in scored) turns_passed = sum(_score(case, "turns_passed") for case in scored) summary: dict[str, float | int] = { "conversations": len(scored), @@ -475,9 +478,9 @@ def _live_summary(report_cases, report_failures=()) -> dict[str, float | int] | _score(case, "cited_map") for case in cited ) / len(cited) - true_refusals = sum(_score(case, "true_refusals") or 0 for case in scored) - false_refusals = sum(_score(case, "false_refusals") or 0 for case in scored) - unanswerable = sum(_score(case, "unanswerable_turns") or 0 for case in scored) + true_refusals = sum(_score(case, "true_refusals") for case in scored) + false_refusals = sum(_score(case, "false_refusals") for case in scored) + unanswerable = sum(_score(case, "unanswerable_turns") for case in scored) refusals = true_refusals + false_refusals summary["unanswerable_turns"] = unanswerable summary["refusals"] = refusals @@ -497,7 +500,7 @@ def _refusal_metrics(report_cases) -> tuple[float, float, int, int] | None: for case in report_cases: refused = case.assertions.get("refused") label = (case.metadata or {}).get("answerability") - if refused is None or label not in ("ANSWERABLE", "UNANSWERABLE"): + if refused is None or label not in REFUSAL_ELIGIBLE_LABELS: continue outcomes.append((label, bool(refused.value))) if not outcomes: @@ -520,18 +523,29 @@ def _filter_qa_corpus(corpus, case_ids: set[str] | None): return corpus.filter(lambda row: row.get("id") in case_ids) -async def run_qa_benchmark( +class _QARun(NamedTuple): + cases: list[Case[Any, Any, dict[str, Any]]] + db: Path + judge_config: ModelConfig + eval_name: str + experiment_metadata: dict[str, Any] + capability_factory: CapabilityFactory + capability_model: Any + + +def _prepare_qa_run( spec: DatasetSpec, config: AppConfig, - limit: int | None = None, - name: str | None = None, - db_path: Path | None = None, - judge_model: ModelConfig | None = None, - target: Target = "rag-capability", - capability_model: ModelConfig | None = None, - case_ids: set[str] | None = None, - document_filter: str | None = None, -) -> ReportCaseFailure[str, str, dict[str, str]] | None: + limit: int | None, + name: str | None, + db_path: Path | None, + judge_model: ModelConfig | None, + target: Target, + capability_model: ModelConfig | None, + case_ids: set[str] | None, + document_filter: str | None, +) -> _QARun: + """Shared setup for the QA runners: cases, models, name and metadata.""" corpus = spec.qa_loader() corpus = _filter_qa_corpus(corpus, case_ids) if limit is not None: @@ -544,7 +558,74 @@ async def run_qa_benchmark( judge_config = judge_model or DEFAULT_JUDGE_MODEL capability_config = _resolve_capability_config(target, config, capability_model) - db = spec.db_path(db_path) + + eval_name = name if name is not None else f"{spec.key}_qa_evaluation" + experiment_metadata = build_experiment_metadata( + dataset_key=spec.key, + test_cases=len(cases), + config=config, + judge_config=judge_config, + target=target, + capability_config=capability_config, + document_filter=document_filter, + ) + experiment_metadata.update(spec.experiment_metadata or {}) + + return _QARun( + cases=cases, + db=spec.db_path(db_path), + judge_config=judge_config, + eval_name=eval_name, + experiment_metadata=experiment_metadata, + capability_factory=_capability_factory_for_target(target), + capability_model=get_model(capability_config, config), + ) + + +def _print_mean_task_time(report_cases, unit: str = "case") -> None: + if not report_cases: + return + mean = sum(case.task_duration for case in report_cases) / len(report_cases) + console.print(f"Avg task time per {unit}: {mean:.2f}s") + + +def _print_failures(failures, show_question: bool = False) -> None: + if not failures: + return + console.print("[red]\nSummary of failures:[/red]") + for failure in failures: + console.print(f"Case: {failure.name}") + if show_question: + console.print(f"Question: {failure.inputs}") + console.print(f"Error: {failure.error_message}") + console.print("") + + +async def run_qa_benchmark( + spec: DatasetSpec, + config: AppConfig, + limit: int | None = None, + name: str | None = None, + db_path: Path | None = None, + judge_model: ModelConfig | None = None, + target: Target = "rag-capability", + capability_model: ModelConfig | None = None, + case_ids: set[str] | None = None, + document_filter: str | None = None, +) -> ReportCaseFailure[str, str, dict[str, str]] | None: + run = _prepare_qa_run( + spec, + config, + limit, + name, + db_path, + judge_model, + target, + capability_model, + case_ids, + document_filter, + ) + cases, judge_config = run.cases, run.judge_config _attach_relevant_uris(cases, spec, limit) citation_evaluator = spec.citation_evaluator @@ -568,43 +649,20 @@ async def run_qa_benchmark( ] if citation_evaluator is not None: evaluators.append(citation_evaluator) - if spec.evaluate_refusal: - evaluators.append( - RefusalJudge( - rubric=REFUSAL_RUBRIC, - model=get_model(judge_config, config), - assertion={"evaluation_name": "refused", "include_reason": False}, - ) + # RefusalJudge scores only cases whose metadata carries an answerability + # label; on unlabeled datasets it returns no score without a judge call. + evaluators.append( + RefusalJudge( + rubric=REFUSAL_RUBRIC, + model=get_model(judge_config, config), + assertion={"evaluation_name": "refused", "include_reason": False}, ) + ) evaluation_dataset = EvalDataset[Any, str, dict[str, Any]]( name=spec.key, cases=cases, evaluators=evaluators ) - eval_name = name if name is not None else f"{spec.key}_qa_evaluation" - experiment_metadata = build_experiment_metadata( - dataset_key=spec.key, - test_cases=len(cases), - config=config, - judge_config=judge_config, - target=target, - capability_config=capability_config, - document_filter=document_filter, - ) - experiment_metadata.update(spec.experiment_metadata or {}) - - async def _evaluate(answer_fn: Callable[[Any], Awaitable[str]]): - return await evaluation_dataset.evaluate( - answer_fn, - name=eval_name, - max_concurrency=1, - progress=True, - metadata=experiment_metadata, - ) - - capability_factory = _capability_factory_for_target(target) - resolved_capability_model = get_model(capability_config, config) - async def answer_question(inputs: str | ConversationInput) -> str: if isinstance(inputs, ConversationInput): question = inputs.question @@ -613,11 +671,11 @@ async def run_qa_benchmark( question = inputs message_history = None result = await run_capability_question( - capability_factory=capability_factory, - db_path=db, + capability_factory=run.capability_factory, + db_path=run.db, config=config, question=question, - capability_model=resolved_capability_model, + capability_model=run.capability_model, document_filter=document_filter, message_history=message_history, ) @@ -633,7 +691,13 @@ async def run_qa_benchmark( set_eval_attribute("citation_status", result.citation_status) return result.answer - report = await _evaluate(answer_question) + report = await evaluation_dataset.evaluate( + answer_question, + name=run.eval_name, + max_concurrency=1, + progress=True, + metadata=run.experiment_metadata, + ) total_processed = len(report.cases) failures = report.failures @@ -660,11 +724,7 @@ async def run_qa_benchmark( console.print(f"Total questions: {total_processed}") console.print(f"Correct answers: {passing_cases}") console.print(f"QA Accuracy: {accuracy:.4f} ({accuracy * 100:.2f}%)") - if report.cases: - mean_task_time = sum(case.task_duration for case in report.cases) / len( - report.cases - ) - console.print(f"Avg task time per case: {mean_task_time:.2f}s") + _print_mean_task_time(report.cases) if citation_evaluator is not None: score_key = citation_evaluator.get_default_evaluation_name() @@ -693,26 +753,16 @@ async def run_qa_benchmark( ) console.print(f"Mean citations per case: {mean_citations:.2f}") - if spec.evaluate_refusal: - metrics = _refusal_metrics(report.cases) - if metrics is not None: - precision, recall, unanswerable, refusals = metrics - console.print( - "\n=== Refusal vs answerability labels ===", style="bold cyan" - ) - console.print(f"Refusal precision: {precision:.2%} | recall: {recall:.2%}") - console.print( - f"UNANSWERABLE turns: {unanswerable} | refusals: {refusals} " - "(PARTIAL excluded)" - ) + if (metrics := _refusal_metrics(report.cases)) is not None: + precision, recall, unanswerable, refusals = metrics + console.print("\n=== Refusal vs answerability labels ===", style="bold cyan") + console.print(f"Refusal precision: {precision:.2%} | recall: {recall:.2%}") + console.print( + f"UNANSWERABLE turns: {unanswerable} | refusals: {refusals} " + "(PARTIAL excluded)" + ) - if failures: - console.print("[red]\nSummary of failures:[/red]") - for failure in failures: - console.print(f"Case: {failure.name}") - console.print(f"Question: {failure.inputs}") - console.print(f"Error: {failure.error_message}") - console.print("") + _print_failures(failures, show_question=True) return failures[0] if failures else None @@ -727,58 +777,44 @@ async def run_live_qa_benchmark( target: Target = "rag-capability", capability_model: ModelConfig | None = None, case_ids: set[str] | None = None, + document_filter: str | None = None, ) -> None: """Replay conversations turn by turn through one capability session. One case per conversation; ``limit`` counts conversations. Answers carry forward as real message history, so prior-turn compaction is exercised. """ - corpus = spec.qa_loader() - corpus = _filter_qa_corpus(corpus, case_ids) - if limit is not None: - corpus = corpus.select(range(min(limit, len(corpus)))) - - cases = [ - spec.qa_case_builder(index, cast(Mapping[str, Any], doc)) - for index, doc in enumerate(corpus, start=1) - ] - - judge_config = judge_model or DEFAULT_JUDGE_MODEL - capability_config = _resolve_capability_config(target, config, capability_model) - db = spec.db_path(db_path) + run = _prepare_qa_run( + spec, + config, + limit, + name, + db_path, + judge_model, + target, + capability_model, + case_ids, + document_filter, + ) evaluation_dataset = EvalDataset[Any, Any, dict[str, Any]]( name=spec.key, - cases=cases, + cases=run.cases, evaluators=[ ConversationEvaluator( rubric=ANSWER_EQUIVALENCE_RUBRIC, - model=get_model(judge_config, config), + model=get_model(run.judge_config, config), ) ], ) - eval_name = name if name is not None else f"{spec.key}_qa_evaluation" - experiment_metadata = build_experiment_metadata( - dataset_key=spec.key, - test_cases=len(cases), - config=config, - judge_config=judge_config, - target=target, - capability_config=capability_config, - ) - experiment_metadata.update(spec.experiment_metadata or {}) - - capability_factory = _capability_factory_for_target(target) - resolved_capability_model = get_model(capability_config, config) - async def answer_conversation(questions: list[str]) -> list[str]: results = await run_capability_conversation( - capability_factory=capability_factory, - db_path=db, + capability_factory=run.capability_factory, + db_path=run.db, config=config, questions=list(questions), - capability_model=resolved_capability_model, + capability_model=run.capability_model, compaction=spec.compaction, ) set_eval_attribute("turn_cited_uris", [r.cited_uris for r in results]) @@ -793,10 +829,10 @@ async def run_live_qa_benchmark( report = await evaluation_dataset.evaluate( answer_conversation, - name=eval_name, + name=run.eval_name, max_concurrency=1, progress=True, - metadata=experiment_metadata, + metadata=run.experiment_metadata, ) summary = _live_summary(report.cases, report.failures) @@ -849,12 +885,7 @@ async def run_live_qa_benchmark( f"{per_turn:.2f}s per turn" ) - if report.failures: - console.print("[red]\nSummary of failures:[/red]") - for failure in report.failures: - console.print(f"Case: {failure.name}") - console.print(f"Error: {failure.error_message}") - console.print("") + _print_failures(report.failures) async def evaluate_dataset( diff --git a/evaluations/evaluations/capability_runner.py b/evaluations/evaluations/capability_runner.py index 89e0e334..2ab4cef6 100644 --- a/evaluations/evaluations/capability_runner.py +++ b/evaluations/evaluations/capability_runner.py @@ -211,8 +211,6 @@ async def run_capability_conversation( config: AppConfig, questions: list[str], capability_model: str | Model, - document_filter: str | None = None, - request_limit: int | None = None, compaction: bool = False, ) -> list[CapabilityRunResult]: """Run a conversation's user turns sequentially through one capability. @@ -229,8 +227,8 @@ async def run_capability_conversation( db_path, config, capability_model, - document_filter, - request_limit, + document_filter=None, + request_limit=None, compaction=compaction, ) history: list[ModelMessage] | None = None diff --git a/evaluations/evaluations/config.py b/evaluations/evaluations/config.py index 23faef87..205905d2 100644 --- a/evaluations/evaluations/config.py +++ b/evaluations/evaluations/config.py @@ -60,7 +60,6 @@ DocumentLoader = Callable[[], Dataset] DocumentMapper = Callable[[Mapping[str, Any]], DocumentPayload | None] RetrievalLoader = Callable[[], Dataset] RetrievalMapper = Callable[[Mapping[str, Any]], RetrievalSample | None] -QAInput = str | ConversationInput CaseBuilder = Callable[[int, Mapping[str, Any]], Case[Any, Any, dict[str, Any]]] @@ -80,7 +79,6 @@ class DatasetSpec: document_limit: int | None = None retrieval_limit: int = 5 ingest_batch_size: int | None = None - evaluate_refusal: bool = False live: bool = False compaction: bool = False experiment_metadata: dict[str, Any] | None = None diff --git a/evaluations/evaluations/datasets/mtrag.py b/evaluations/evaluations/datasets/mtrag.py index 420913fd..61cc29b0 100644 --- a/evaluations/evaluations/datasets/mtrag.py +++ b/evaluations/evaluations/datasets/mtrag.py @@ -165,7 +165,7 @@ def _task_to_record( } -def load_clapnq_qa() -> Dataset: +def _qa_records() -> list[dict[str, Any]]: path = _download(_GEN_TASKS_FILE) qrels = _load_qrels() records = [] @@ -175,7 +175,11 @@ def load_clapnq_qa() -> Dataset: record = _task_to_record(json.loads(line), qrels) if record is not None: records.append(record) - return Dataset.from_list(records) + return records + + +def load_clapnq_qa() -> Dataset: + return Dataset.from_list(_qa_records()) def build_mtrag_case( @@ -225,17 +229,15 @@ def _group_conversations(records: list[dict[str, Any]]) -> list[dict[str, Any]]: "answerability": task["answerability"], "multi_turn_type": task["multi_turn_type"], "question_type": list(task["question_type"]), + "relevant_uris": list(task["relevant_uris"] or []), } - if task["relevant_uris"]: - turn["relevant_uris"] = list(task["relevant_uris"]) turns.append(turn) conversations.append({"id": conversation_id, "turns": turns}) return conversations def load_clapnq_conversations() -> Dataset: - corpus = load_clapnq_qa() - return Dataset.from_list(_group_conversations([dict(row) for row in corpus])) + return Dataset.from_list(_group_conversations(_qa_records())) def build_mtrag_live_case( @@ -243,11 +245,7 @@ def build_mtrag_live_case( ) -> Case[list[str], list[str], dict[str, Any]]: questions = [turn["question"] for turn in doc["turns"]] metadata_turns = [ - { - key: value - for key, value in turn.items() - if key != "question" and value is not None - } + {key: value for key, value in turn.items() if key != "question"} for turn in doc["turns"] ] return Case( @@ -277,7 +275,6 @@ def _mtrag_spec(key: str, variant: str) -> DatasetSpec: citation_evaluator=CitationMAPEvaluator(), retrieval_limit=10, ingest_batch_size=512, - evaluate_refusal=True, experiment_metadata={"mtrag_mode": "gold_prefix"}, ) diff --git a/evaluations/evaluations/evaluators/__init__.py b/evaluations/evaluations/evaluators/__init__.py index fd3fe169..97fb771b 100644 --- a/evaluations/evaluations/evaluators/__init__.py +++ b/evaluations/evaluations/evaluators/__init__.py @@ -7,12 +7,17 @@ from evaluations.evaluators.judge import ( ) from evaluations.evaluators.map import MAPEvaluator from evaluations.evaluators.number_match import NumberMatchEvaluator -from evaluations.evaluators.refusal import REFUSAL_RUBRIC, RefusalJudge +from evaluations.evaluators.refusal import ( + REFUSAL_ELIGIBLE_LABELS, + REFUSAL_RUBRIC, + RefusalJudge, +) from evaluations.evaluators.retrieval import NDCGEvaluator, RecallEvaluator from evaluations.evaluators.transcript import TranscriptLLMJudge __all__ = [ "ANSWER_EQUIVALENCE_RUBRIC", + "REFUSAL_ELIGIBLE_LABELS", "REFUSAL_RUBRIC", "CitationMAPEvaluator", "ConversationEvaluator", diff --git a/evaluations/evaluations/evaluators/conversation.py b/evaluations/evaluations/evaluators/conversation.py index 8b97c60f..78ca74e4 100644 --- a/evaluations/evaluations/evaluators/conversation.py +++ b/evaluations/evaluations/evaluators/conversation.py @@ -9,9 +9,7 @@ from pydantic_evals.evaluators.llm_as_a_judge import ( ) from evaluations.evaluators.citation import average_precision -from evaluations.evaluators.refusal import REFUSAL_RUBRIC - -_REFUSAL_LABELS = ("ANSWERABLE", "UNANSWERABLE") +from evaluations.evaluators.refusal import REFUSAL_ELIGIBLE_LABELS, REFUSAL_RUBRIC @dataclass @@ -81,7 +79,7 @@ class ConversationEvaluator(Evaluator): ) label = turn.get("answerability") - if label in _REFUSAL_LABELS: + if label in REFUSAL_ELIGIBLE_LABELS: try: refused = ( await judge_output(answer, REFUSAL_RUBRIC, self.model) diff --git a/evaluations/evaluations/evaluators/map.py b/evaluations/evaluations/evaluators/map.py index 1746f6ad..e4fe5a32 100644 --- a/evaluations/evaluations/evaluators/map.py +++ b/evaluations/evaluations/evaluators/map.py @@ -2,6 +2,8 @@ from dataclasses import dataclass from pydantic_evals.evaluators import Evaluator, EvaluatorContext +from evaluations.evaluators.citation import average_precision + @dataclass class MAPEvaluator(Evaluator): @@ -28,22 +30,6 @@ class MAPEvaluator(Evaluator): if ctx.metadata is None: return 0.0 relevant_uris = set(ctx.metadata.get("relevant_uris", [])) - retrieved_uris = ctx.output - if not relevant_uris: return 0.0 - - num_relevant = len(relevant_uris) - precisions = [] - num_relevant_found = 0 - - for rank, uri in enumerate(retrieved_uris, start=1): - if uri in relevant_uris: - num_relevant_found += 1 - precision_at_k = num_relevant_found / rank - precisions.append(precision_at_k) - - if not precisions: - return 0.0 - - return sum(precisions) / num_relevant + return average_precision(list(ctx.output), relevant_uris) diff --git a/evaluations/evaluations/evaluators/refusal.py b/evaluations/evaluations/evaluators/refusal.py index ede0b4a5..435720fc 100644 --- a/evaluations/evaluations/evaluators/refusal.py +++ b/evaluations/evaluations/evaluators/refusal.py @@ -1,4 +1,5 @@ from dataclasses import dataclass +from typing import Final from pydantic_evals.evaluators import EvaluatorContext, LLMJudge from pydantic_evals.evaluators.evaluator import EvaluatorOutput @@ -9,7 +10,7 @@ REFUSAL_RUBRIC = ( "substantive answer." ) -_ELIGIBLE_LABELS = ("ANSWERABLE", "UNANSWERABLE") +REFUSAL_ELIGIBLE_LABELS: Final = ("ANSWERABLE", "UNANSWERABLE") @dataclass @@ -23,6 +24,6 @@ class RefusalJudge(LLMJudge): async def evaluate(self, ctx: EvaluatorContext) -> EvaluatorOutput: label = (ctx.metadata or {}).get("answerability") - if label not in _ELIGIBLE_LABELS: + if label not in REFUSAL_ELIGIBLE_LABELS: return {} return await super().evaluate(ctx) diff --git a/evaluations/tests/test_benchmark.py b/evaluations/tests/test_benchmark.py index a9aa6454..5345f571 100644 --- a/evaluations/tests/test_benchmark.py +++ b/evaluations/tests/test_benchmark.py @@ -332,7 +332,7 @@ class TestLiveSummary: def test_none_without_scored_cases(self) -> None: from evaluations.benchmark import _live_summary - assert _live_summary([self._case({})]) is None + assert _live_summary([self._case({})], []) is None def test_micro_rate_uses_judged_turns(self) -> None: from evaluations.benchmark import _live_summary @@ -352,7 +352,7 @@ class TestLiveSummary: ) ] - summary = _live_summary(cases) + summary = _live_summary(cases, []) assert summary is not None assert summary["micro_pass_rate"] == 1.0 @@ -784,8 +784,9 @@ class TestRunQaBenchmarkCapabilityTarget: # (the capability manages its own client via lifespan). mock_haiku.assert_not_called() # capability model defaults to qa.model when not provided - capability_call = mock_get_model.call_args_list[-1] - assert capability_call[0][0] == AppConfig().qa.model + assert any( + call[0][0] == AppConfig().qa.model for call in mock_get_model.call_args_list + ) assert mock_run_capability is capability_run @pytest.mark.asyncio diff --git a/evaluations/tests/test_mtrag.py b/evaluations/tests/test_mtrag.py index 0304a2fa..90a7cc7b 100644 --- a/evaluations/tests/test_mtrag.py +++ b/evaluations/tests/test_mtrag.py @@ -139,9 +139,6 @@ class TestSpecs: } assert isinstance(spec.citation_evaluator, CitationMAPEvaluator) - def test_refusal_evaluation_enabled(self) -> None: - assert MTRAG_CLAPNQ_SPEC.evaluate_refusal is True - class TestGenerationTasks: def test_task_to_record(self) -> None: @@ -255,7 +252,7 @@ class TestLiveConversations: } other_case = build_mtrag_live_case(2, conversations[1]) assert other_case.metadata is not None - assert "relevant_uris" not in other_case.metadata["turns"][0] + assert other_case.metadata["turns"][0]["relevant_uris"] == [] def test_live_spec(self) -> None: assert DATASETS["mtrag_clapnq_live"] is MTRAG_CLAPNQ_LIVE_SPEC diff --git a/haiku_rag_slim/haiku/rag/client/documents.py b/haiku_rag_slim/haiku/rag/client/documents.py index 1f5c4aae..7257fcbd 100644 --- a/haiku_rag_slim/haiku/rag/client/documents.py +++ b/haiku_rag_slim/haiku/rag/client/documents.py @@ -299,21 +299,16 @@ async def _store_documents_with_chunks( Embeds any chunks that lack embeddings, then writes the documents, chunks, and document_items tables once apiece. Restores all tables on any failure. """ - missing = [ - chunk - for _, chunks, _ in prepared - for chunk in chunks - if chunk.embedding is None - ] - if missing: - from haiku.rag.embeddings import embed_chunks - - embedded_flat = await embed_chunks(missing, client.embedder, client._config) - # Assign positionally: duplicate chunk texts across documents make a - # content-keyed lookup ambiguous. - for chunk, with_embedding in zip(missing, embedded_flat): - chunk.embedding = with_embedding.embedding - embedded: list[list[Chunk]] = [chunks for _, chunks, _ in prepared] + flat = await ensure_chunks_embedded( + client._config, + [chunk for _, chunks, _ in prepared for chunk in chunks], + client.embedder, + ) + embedded: list[list[Chunk]] = [] + position = 0 + for _, chunks, _ in prepared: + embedded.append(flat[position : position + len(chunks)]) + position += len(chunks) def _extract_all_items(): return [extract_items("", d) for _, _, d in prepared] diff --git a/haiku_rag_slim/haiku/rag/client/processing.py b/haiku_rag_slim/haiku/rag/client/processing.py index 87f59ff3..4d541658 100644 --- a/haiku_rag_slim/haiku/rag/client/processing.py +++ b/haiku_rag_slim/haiku/rag/client/processing.py @@ -352,16 +352,10 @@ async def ensure_chunks_embedded( embedded = await embed_chunks(chunks_to_embed, embedder, config) - # Build result maintaining original order - embedded_map = {(c.content, c.order): c for c in embedded} - result = [] - for ch in chunks: - if ch.embedding is not None: - result.append(ch) - else: - result.append(embedded_map[(ch.content, ch.order)]) - - return result + # embed_chunks preserves input order; fill positionally, since duplicate + # chunk texts across documents make a content-keyed lookup ambiguous. + filled = iter(embedded) + return [ch if ch.embedding is not None else next(filled) for ch in chunks] def get_extension_from_content_type_or_url(url: str, content_type: str) -> str: