diff --git a/litellm/__init__.py b/litellm/__init__.py index 1784deeb797..410402beb6e 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -1453,6 +1453,7 @@ from .batches.main import * from .images.main import * from .videos.main import * from .batch_completion.main import * +from .fusion.main import fusion, afusion, FusionStrategy from .rerank_api.main import * from .messages.dispatch import * from .responses.dispatch import * diff --git a/litellm/fusion/__init__.py b/litellm/fusion/__init__.py new file mode 100644 index 00000000000..1b877b9e2a1 --- /dev/null +++ b/litellm/fusion/__init__.py @@ -0,0 +1,3 @@ +from .main import FusionStrategy, afusion, fusion + +__all__ = ["FusionStrategy", "afusion", "fusion"] # mutable-ok: module __all__ list, Python convention diff --git a/litellm/fusion/main.py b/litellm/fusion/main.py new file mode 100644 index 00000000000..9b736c1e709 --- /dev/null +++ b/litellm/fusion/main.py @@ -0,0 +1,444 @@ +""" +litellm.fusion — Multi-model Fusion with optional judge synthesis. + +Sends the same prompt to N models in parallel. Without a judge model, returns +all panel responses as a list. With a judge model, synthesizes a final answer +and includes panel responses by default (set include_panel=False to suppress). + +Inspired by OpenRouter Fusion / Sakana Fugu / Self-Consistency (Wang et al., 2022). + +Usage: + # Panel only — get all responses, pick yourself + responses = litellm.fusion( + models=["gpt-4o", "claude-3-5-sonnet", "gemini-2.0-flash"], + messages=[{"role": "user", "content": "Explain quantum entanglement"}], + ) + # returns: list[ModelResponse] + + # Judge synthesis — synthesized answer + panel included by default + response = litellm.fusion( + models=["gpt-4o", "claude-3-5-sonnet", "gemini-2.0-flash"], + judge_model="gpt-4o", + messages=[{"role": "user", "content": "Explain quantum entanglement"}], + ) + # returns: ModelResponse + # response._hidden_params["fusion"]["panel_responses"] <- individual responses + + # Judge synthesis — panel excluded + response = litellm.fusion( + models=["gpt-4o", "claude-3-5-sonnet", "gemini-2.0-flash"], + judge_model="gpt-4o", + messages=[{"role": "user", "content": "Explain quantum entanglement"}], + include_panel=False, + ) + # returns: ModelResponse (no panel_responses in metadata) + + # Async variants + responses = await litellm.afusion(models=[...], messages=[...]) + response = await litellm.afusion(models=[...], judge_model="...", messages=[...]) +""" + +from __future__ import annotations + +import asyncio +import time +import uuid +from collections.abc import Sequence +from typing import Final, Literal, TypeAlias, overload + +import litellm +from litellm.types.utils import ModelResponse, Usage + +# --------------------------------------------------------------------------- +# Judge prompt templates +# --------------------------------------------------------------------------- + +_JUDGE_PROMPT_SINGLE: Final = ( # rebind-ok: module-level constant, never rebound + "You received the following responses from {n} AI models for this user request:\n\n" + "{responses}\n\n" + "Your task: synthesize a single, comprehensive, and accurate final answer.\n" + "- Identify the points of agreement and the best reasoning across responses.\n" + "- Resolve contradictions by selecting the most well-supported position.\n" + "- Fill in gaps where some models provided information others missed.\n" + "- Write the final answer as if you generated it directly" + " — no meta-commentary about the synthesis process.\n" +) + +_JUDGE_PROMPT_MAJORITY: Final = ( # rebind-ok: module-level constant, never rebound + "You received the following responses from {n} AI models for this user request:\n\n" + "{responses}\n\n" + "Select the single best response. Output ONLY the text of the chosen response, verbatim.\n" +) + +_JUDGE_PROMPT_BEST_OF_N: Final = ( # rebind-ok: module-level constant, never rebound + "You received the following responses from {n} AI models for this user request:\n\n" + "{responses}\n\n" + "Score each response on a scale of 1-10 for accuracy, completeness, and clarity.\n" + "Then output the highest-scored response verbatim.\n" +) + +_JUDGE_PROMPTS: Final = { # mutable-ok: module-level lookup table, never replaced + "single_judge": _JUDGE_PROMPT_SINGLE, + "majority_vote": _JUDGE_PROMPT_MAJORITY, + "best_of_n": _JUDGE_PROMPT_BEST_OF_N, +} + +FusionStrategy: TypeAlias = Literal["single_judge", "majority_vote", "best_of_n"] + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _build_judge_messages( + original_messages: Sequence[dict], # mutable-ok: OpenAI message dicts, caller-owned + panel_responses: Sequence[ModelResponse], # mutable-ok: panel results, caller-owned + panel_models: Sequence[str], # mutable-ok: model names, caller-owned + strategy: FusionStrategy, +) -> list[dict]: # mutable-ok: message list consumed directly by litellm.acompletion + """Build the message list for the judge call.""" + user_request: Final = next( + ( + m["content"] + for m in reversed(list(original_messages)) # mutable-ok: ephemeral list for reversed(), never stored + if m.get("role") == "user" + ), + "", + ) + + response_blocks: list[str] = [] # mutable-ok: accumulated then joined, never escapes # rebind-ok: never rebound + for i, (model, resp) in enumerate(zip(panel_models, panel_responses), start=1): + content = "" + if resp and resp.choices: + content = resp.choices[0].message.content or "" + response_blocks.append(f"[Response {i} - {model}]\n{content}") + + responses_text: Final = "\n\n".join(response_blocks) + template: Final = _JUDGE_PROMPTS[strategy] + judge_system: Final = template.format(n=len(panel_responses), responses=responses_text) + + messages: list[dict] = [] # mutable-ok: accumulated then returned immediately # rebind-ok: never rebound + + # Preserve any existing system prompt + for m in original_messages: + if m.get("role") == "system": + messages.append(m) + break + + messages.append({"role": "system", "content": judge_system}) # mutable-ok: dict consumed by litellm + messages.append({"role": "user", "content": user_request}) # mutable-ok: dict consumed by litellm + return messages + + +def _sum_usage(responses: Sequence[ModelResponse]) -> Usage: + """Sum token usage across a list of ModelResponse objects.""" + total_prompt = 0 # rebind-ok: accumulator, augmented in loop + total_completion = 0 # rebind-ok: accumulator, augmented in loop + total_tokens = 0 # rebind-ok: accumulator, augmented in loop + for r in responses: + if r and r.usage: + total_prompt += r.usage.prompt_tokens or 0 + total_completion += r.usage.completion_tokens or 0 + total_tokens += r.usage.total_tokens or 0 + return Usage( + prompt_tokens=total_prompt, + completion_tokens=total_completion, + total_tokens=total_tokens, + ) + + +def _merge_usage( + panel_responses: Sequence[ModelResponse], + judge_response: ModelResponse, +) -> Usage: + """Sum token usage across all panel calls + judge call.""" + return _sum_usage( + (*panel_responses, judge_response) + ) # mutable-ok: ephemeral tuple, passed immediately to _sum_usage + + +def _build_fusion_response( + judge_response: ModelResponse, + panel_responses: Sequence[ModelResponse], # mutable-ok: stored in metadata, caller-owned + panel_models: Sequence[str], # mutable-ok: stored in metadata, caller-owned + judge_model: str, + include_panel: bool, + original_model_tag: str, +) -> ModelResponse: + """Wrap judge response in a ModelResponse tagged with fusion metadata.""" + merged_usage: Final = _merge_usage(panel_responses, judge_response) + + response: Final = ModelResponse( + id=f"fusion-{uuid.uuid4().hex}", + choices=judge_response.choices, + created=int(time.time()), + model=original_model_tag, + usage=merged_usage, + object="chat.completion", + ) + + fusion_meta: dict = { # mutable-ok: conditionally extended before being stored # rebind-ok: never rebound + "panel_models": panel_models, + "judge_model": judge_model, + } + if include_panel: + fusion_meta["panel_responses"] = panel_responses # mutable-ok: stored in _hidden_params + + response._hidden_params = {"fusion": fusion_meta} # mutable-ok: ModelResponse._hidden_params is the SDK extension point # fmt: skip + return response + + +async def _run_panel( + models: Sequence[str], # mutable-ok: iterated read-only + messages: Sequence[dict], # mutable-ok: forwarded to litellm.acompletion + panel_kwargs: dict, # mutable-ok: forwarded kwargs dict, caller-owned +) -> tuple[list[ModelResponse], list[str]]: # mutable-ok: results consumed immediately by afusion + """Fan out to all panel models in parallel; filter failed calls.""" + tasks: Final = [ # mutable-ok: list of coroutines passed to asyncio.gather + litellm.acompletion( + model=m, + messages=list(messages), # mutable-ok: ephemeral list copy for acompletion, never stored + stream=False, + **panel_kwargs, # kwargs-ok: forwarded to litellm.acompletion, varies per call + ) + for m in models + ] + raw: Final = list(await asyncio.gather(*tasks, return_exceptions=True)) # mutable-ok: gather results consumed below + + valid_responses: list[ModelResponse] = [] # mutable-ok: accumulated then returned # rebind-ok: never rebound + valid_models: list[str] = [] # mutable-ok: accumulated then returned # rebind-ok: never rebound + for model, resp in zip(models, raw): + if isinstance(resp, Exception): + litellm.utils.print_verbose(f"fusion: panel model {model!r} failed: {resp}") + else: + valid_responses.append(resp) + valid_models.append(model) + + if not valid_responses: + raise RuntimeError(f"fusion: all panel models failed. Errors: {raw}") + + return valid_responses, valid_models + + +# --------------------------------------------------------------------------- +# Overloads for precise return-type inference +# --------------------------------------------------------------------------- + + +@overload +async def afusion( + models: Sequence[str], # mutable-ok: OpenAI-compatible, matches litellm API + messages: Sequence[dict], # mutable-ok: OpenAI-compatible message list + *, + judge_model: None = ..., + strategy: FusionStrategy = ..., + include_panel: bool = ..., + timeout: float | None = ..., + temperature: float | None = ..., + max_tokens: int | None = ..., + **kwargs: object, # kwargs-ok: forwarded to litellm.acompletion, varies per provider +) -> list[ModelResponse]: ... # mutable-ok: panel responses list + + +@overload +async def afusion( + models: Sequence[str], # mutable-ok: OpenAI-compatible, matches litellm API + messages: Sequence[dict], # mutable-ok: OpenAI-compatible message list + *, + judge_model: str, + strategy: FusionStrategy = ..., + include_panel: bool = ..., + timeout: float | None = ..., + temperature: float | None = ..., + max_tokens: int | None = ..., + **kwargs: object, # kwargs-ok: forwarded to litellm.acompletion, varies per provider +) -> ModelResponse: ... + + +# --------------------------------------------------------------------------- +# Core async implementation +# --------------------------------------------------------------------------- + + +async def afusion( + models: Sequence[str], # mutable-ok: OpenAI-compatible, matches litellm API + messages: Sequence[dict], # mutable-ok: OpenAI-compatible message list + *, + judge_model: str | None = None, + strategy: FusionStrategy = "single_judge", + include_panel: bool = True, + timeout: float | None = None, + temperature: float | None = None, + max_tokens: int | None = None, + **kwargs: object, # kwargs-ok: forwarded to litellm.acompletion, varies per provider +) -> list[ModelResponse] | ModelResponse: # mutable-ok: list[ModelResponse] is the panel-only branch + """ + Async fusion: call panel models in parallel, optionally synthesize with judge. + + Args: + models: List of panel model identifiers (>= 2 recommended). + messages: OpenAI-compatible message list. + judge_model: If provided, synthesizes a final answer from panel responses. + If omitted, returns the raw list of panel responses. + strategy: How the judge processes responses (only used when judge_model is set). + - "single_judge": synthesize a new combined answer (default) + - "majority_vote": pick the single best response verbatim + - "best_of_n": score and select highest-scored response + include_panel: When judge_model is set, whether to include individual panel + responses in response._hidden_params["fusion"]["panel_responses"]. + Defaults to True. Ignored when judge_model is None. + timeout: Per-call timeout in seconds (applied to panel + judge calls). + temperature: Forwarded to panel models. Judge always uses temperature=0. + max_tokens: Forwarded to panel models. + **kwargs: Any additional litellm.acompletion kwargs forwarded to panel calls. + + Returns: + - list[ModelResponse] when judge_model is None + - ModelResponse (synthesized) when judge_model is provided + """ + if not models: + raise ValueError("fusion: `models` must be a non-empty list") + if strategy not in _JUDGE_PROMPTS: + raise ValueError( + f"fusion: unknown strategy {strategy!r}. Choose from {list(_JUDGE_PROMPTS)}" # mutable-ok: ephemeral list for error message + ) + + panel_kwargs: dict = dict(kwargs) # mutable-ok: working copy of kwargs, extended below # rebind-ok: never rebound + if temperature is not None: + panel_kwargs["temperature"] = temperature + if max_tokens is not None: + panel_kwargs["max_tokens"] = max_tokens + if timeout is not None: + panel_kwargs["timeout"] = timeout + + # 1. Fan out to all panel models in parallel + valid_responses, valid_models = await _run_panel( + list(models), # mutable-ok: ephemeral list, converts Sequence for _run_panel + list(messages), # mutable-ok: ephemeral list, converts Sequence for _run_panel + panel_kwargs, + ) + + # 2. No judge — return panel responses as-is + if judge_model is None: + return valid_responses + + # 3. Judge synthesis + judge_messages: Final = _build_judge_messages( + original_messages=messages, + panel_responses=valid_responses, + panel_models=valid_models, + strategy=strategy, + ) + + judge_kwargs: dict = {} # mutable-ok: conditionally extended before use # rebind-ok: never rebound + if timeout is not None: + judge_kwargs["timeout"] = timeout + + judge_response: Final = await litellm.acompletion( + model=judge_model, + messages=judge_messages, + stream=False, + temperature=0, # deterministic synthesis + **judge_kwargs, # kwargs-ok: forwarded to litellm.acompletion + ) + + # 4. Wrap into a single ModelResponse with merged metadata + return _build_fusion_response( + judge_response=judge_response, + panel_responses=valid_responses, + panel_models=valid_models, + judge_model=judge_model, + include_panel=include_panel, + original_model_tag=f"fusion/{'+'.join(valid_models)}", + ) + + +# --------------------------------------------------------------------------- +# Sync overloads +# --------------------------------------------------------------------------- + + +@overload +def fusion( + models: Sequence[str], # mutable-ok: OpenAI-compatible, matches litellm API + messages: Sequence[dict], # mutable-ok: OpenAI-compatible message list + *, + judge_model: None = ..., + strategy: FusionStrategy = ..., + include_panel: bool = ..., + timeout: float | None = ..., + temperature: float | None = ..., + max_tokens: int | None = ..., + **kwargs: object, # kwargs-ok: forwarded to litellm.acompletion, varies per provider +) -> list[ModelResponse]: ... # mutable-ok: panel responses list + + +@overload +def fusion( + models: Sequence[str], # mutable-ok: OpenAI-compatible, matches litellm API + messages: Sequence[dict], # mutable-ok: OpenAI-compatible message list + *, + judge_model: str, + strategy: FusionStrategy = ..., + include_panel: bool = ..., + timeout: float | None = ..., + temperature: float | None = ..., + max_tokens: int | None = ..., + **kwargs: object, # kwargs-ok: forwarded to litellm.acompletion, varies per provider +) -> ModelResponse: ... + + +def fusion( + models: Sequence[str], # mutable-ok: OpenAI-compatible, matches litellm API + messages: Sequence[dict], # mutable-ok: OpenAI-compatible message list + *, + judge_model: str | None = None, + strategy: FusionStrategy = "single_judge", + include_panel: bool = True, + timeout: float | None = None, + temperature: float | None = None, + max_tokens: int | None = None, + **kwargs: object, # kwargs-ok: forwarded to litellm.acompletion, varies per provider +) -> list[ModelResponse] | ModelResponse: # mutable-ok: list[ModelResponse] is the panel-only branch + """ + Sync fusion: call panel models in parallel, optionally synthesize with judge. + + Thin sync wrapper around :func:`afusion`. For async contexts, prefer + :func:`afusion` directly. + + Args: + models: List of panel model identifiers (>= 2 recommended). + messages: OpenAI-compatible message list. + judge_model: If provided, synthesizes a final answer from panel responses. + If omitted, returns the raw list of panel responses. + strategy: Synthesis strategy ("single_judge", "majority_vote", "best_of_n"). + Only used when judge_model is provided. + include_panel: When judge_model is set, whether to attach individual panel + responses to response._hidden_params["fusion"]["panel_responses"]. + Defaults to True. Ignored when judge_model is None. + timeout: Per-call timeout in seconds. + temperature: Panel model temperature (judge always uses 0). + max_tokens: Forwarded to panel models. + **kwargs: Any additional litellm.completion kwargs. + + Returns: + - list[ModelResponse] when judge_model is None + - ModelResponse (synthesized) when judge_model is provided + """ + return asyncio.run( + afusion( + models=models, + messages=messages, + judge_model=judge_model, + strategy=strategy, + include_panel=include_panel, + timeout=timeout, + temperature=temperature, + max_tokens=max_tokens, + **kwargs, # kwargs-ok: forwarded to afusion + ) + ) + + +__all__ = ["FusionStrategy", "afusion", "fusion"] # mutable-ok: module __all__, Python convention diff --git a/tests/test_litellm/test_fusion.py b/tests/test_litellm/test_fusion.py new file mode 100644 index 00000000000..e36e5b65659 --- /dev/null +++ b/tests/test_litellm/test_fusion.py @@ -0,0 +1,428 @@ +""" +Tests for litellm.fusion and litellm.afusion. + +Uses unittest.mock to avoid real API calls. +""" +from __future__ import annotations + +from unittest.mock import patch + +import pytest + +import litellm +from litellm.fusion.main import ( + FusionStrategy, + _build_judge_messages, + _merge_usage, + _sum_usage, + afusion, + fusion, +) +from litellm.types.utils import Choices, Message, ModelResponse, Usage + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _make_response(content: str, model: str = "gpt-4o", prompt_tokens: int = 10) -> ModelResponse: + resp = ModelResponse( + id=f"mock-{model}", + choices=[ + Choices( + index=0, + message=Message(role="assistant", content=content), + finish_reason="stop", + ) + ], + model=model, + usage=Usage( + prompt_tokens=prompt_tokens, + completion_tokens=len(content.split()), + total_tokens=prompt_tokens + len(content.split()), + ), + ) + return resp + + +MESSAGES = [{"role": "user", "content": "What is 2+2?"}] +PANEL_MODELS = ["gpt-4o", "claude-3-5-sonnet", "gemini-2.0-flash"] +JUDGE_MODEL = "gpt-4o-judge" # distinct from panel to avoid index confusion + + +def _panel_mock(panel_responses, judge_response=None): + """Return an async mock that routes by model name.""" + async def mock_acompletion(**kwargs): + m = kwargs["model"] + if judge_response is not None and m == JUDGE_MODEL: + return judge_response + if m in PANEL_MODELS: + return panel_responses[PANEL_MODELS.index(m)] + return panel_responses[0] + return mock_acompletion + + +# --------------------------------------------------------------------------- +# Unit: _build_judge_messages +# --------------------------------------------------------------------------- + +class TestBuildJudgeMessages: + def _panel(self): + return [ + _make_response("The answer is 4.", "gpt-4o"), + _make_response("4", "claude-3-5-sonnet"), + ] + + def test_contains_user_question(self): + msgs = _build_judge_messages(MESSAGES, self._panel(), PANEL_MODELS[:2], "single_judge") + combined = " ".join(m["content"] for m in msgs) + assert "2+2" in combined + + def test_contains_panel_responses(self): + msgs = _build_judge_messages(MESSAGES, self._panel(), PANEL_MODELS[:2], "single_judge") + combined = " ".join(m["content"] for m in msgs) + assert "The answer is 4." in combined + assert "4" in combined + + def test_strategy_majority_vote(self): + msgs = _build_judge_messages(MESSAGES, self._panel(), PANEL_MODELS[:2], "majority_vote") + system_msg = next(m for m in msgs if m["role"] == "system") + assert "Select the single best response" in system_msg["content"] + + def test_strategy_best_of_n(self): + msgs = _build_judge_messages(MESSAGES, self._panel(), PANEL_MODELS[:2], "best_of_n") + system_msg = next(m for m in msgs if m["role"] == "system") + assert "Score each response" in system_msg["content"] + + def test_ends_with_user_message(self): + msgs = _build_judge_messages(MESSAGES, self._panel(), PANEL_MODELS[:2], "single_judge") + assert msgs[-1]["role"] == "user" + + +# --------------------------------------------------------------------------- +# Unit: _sum_usage / _merge_usage +# --------------------------------------------------------------------------- + +class TestUsageHelpers: + def test_sum_usage(self): + responses = [ + _make_response("hello", prompt_tokens=10), # completion=1 + _make_response("world more words", prompt_tokens=20), # completion=3 + ] + u = _sum_usage(responses) + assert u.prompt_tokens == 30 + assert u.completion_tokens == 4 + + def test_merge_usage_sums_panel_and_judge(self): + panel = [ + _make_response("hello", prompt_tokens=10), + _make_response("world more words", prompt_tokens=20), + ] + judge = _make_response("final answer here today", prompt_tokens=5) + merged = _merge_usage(panel, judge) + assert merged.prompt_tokens == 35 # 10+20+5 + assert merged.completion_tokens == 1 + 3 + 4 + + def test_handles_none_usage(self): + panel = [_make_response("hi")] + panel[0].usage = None + judge = _make_response("answer", prompt_tokens=5) + merged = _merge_usage(panel, judge) + assert merged.prompt_tokens == 5 + + +# --------------------------------------------------------------------------- +# afusion — panel only (no judge) +# --------------------------------------------------------------------------- + +class TestAFusionPanelOnly: + @pytest.fixture + def panel_responses(self): + return [_make_response(f"Answer {m}", m) for m in PANEL_MODELS] + + @pytest.mark.asyncio + async def test_returns_list_of_model_responses(self, panel_responses): + with patch("litellm.acompletion", side_effect=_panel_mock(panel_responses)): + result = await afusion(models=PANEL_MODELS, messages=MESSAGES) + + assert isinstance(result, list) + assert len(result) == len(PANEL_MODELS) + assert all(isinstance(r, ModelResponse) for r in result) + + @pytest.mark.asyncio + async def test_returns_all_panel_contents(self, panel_responses): + with patch("litellm.acompletion", side_effect=_panel_mock(panel_responses)): + result = await afusion(models=PANEL_MODELS, messages=MESSAGES) + + contents = {r.choices[0].message.content for r in result} + assert contents == {f"Answer {m}" for m in PANEL_MODELS} + + @pytest.mark.asyncio + async def test_no_judge_call_when_judge_model_omitted(self, panel_responses): + called_models: list[str] = [] + + async def mock_acompletion(**kwargs): + called_models.append(kwargs["model"]) + return panel_responses[PANEL_MODELS.index(kwargs["model"])] + + with patch("litellm.acompletion", side_effect=mock_acompletion): + await afusion(models=PANEL_MODELS, messages=MESSAGES) + + # Exactly the panel models, no judge + assert sorted(called_models) == sorted(PANEL_MODELS) + + @pytest.mark.asyncio + async def test_partial_failure_returns_remaining(self, panel_responses): + async def mock_acompletion(**kwargs): + if kwargs["model"] == "gemini-2.0-flash": + raise RuntimeError("API error") + return panel_responses[PANEL_MODELS.index(kwargs["model"])] + + with patch("litellm.acompletion", side_effect=mock_acompletion): + result = await afusion(models=PANEL_MODELS, messages=MESSAGES) + + assert isinstance(result, list) + assert len(result) == 2 # gemini dropped + + +# --------------------------------------------------------------------------- +# afusion — with judge +# --------------------------------------------------------------------------- + +class TestAFusionWithJudge: + @pytest.fixture + def panel_responses(self): + return [_make_response(f"Answer {m}", m) for m in PANEL_MODELS] + + @pytest.fixture + def judge_response(self): + return _make_response("Synthesized final answer", JUDGE_MODEL) + + @pytest.mark.asyncio + async def test_returns_single_model_response(self, panel_responses, judge_response): + with patch("litellm.acompletion", side_effect=_panel_mock(panel_responses, judge_response)): + result = await afusion(models=PANEL_MODELS, judge_model=JUDGE_MODEL, messages=MESSAGES) + + assert isinstance(result, ModelResponse) + assert result.choices[0].message.content == "Synthesized final answer" + + @pytest.mark.asyncio + async def test_calls_all_panel_and_judge(self, panel_responses, judge_response): + called: list[str] = [] + + async def mock_acompletion(**kwargs): + called.append(kwargs["model"]) + return judge_response if kwargs["model"] == JUDGE_MODEL else panel_responses[PANEL_MODELS.index(kwargs["model"])] + + with patch("litellm.acompletion", side_effect=mock_acompletion): + await afusion(models=PANEL_MODELS, judge_model=JUDGE_MODEL, messages=MESSAGES) + + assert set(PANEL_MODELS).issubset(set(called)) + assert JUDGE_MODEL in called + + @pytest.mark.asyncio + async def test_include_panel_true_by_default(self, panel_responses, judge_response): + with patch("litellm.acompletion", side_effect=_panel_mock(panel_responses, judge_response)): + result = await afusion(models=PANEL_MODELS, judge_model=JUDGE_MODEL, messages=MESSAGES) + + fusion_meta = result._hidden_params["fusion"] + assert "panel_responses" in fusion_meta + assert len(fusion_meta["panel_responses"]) == len(PANEL_MODELS) + assert all(isinstance(r, ModelResponse) for r in fusion_meta["panel_responses"]) + + @pytest.mark.asyncio + async def test_include_panel_false_excludes_panel(self, panel_responses, judge_response): + with patch("litellm.acompletion", side_effect=_panel_mock(panel_responses, judge_response)): + result = await afusion( + models=PANEL_MODELS, + judge_model=JUDGE_MODEL, + messages=MESSAGES, + include_panel=False, + ) + + fusion_meta = result._hidden_params["fusion"] + assert "panel_responses" not in fusion_meta + + @pytest.mark.asyncio + async def test_fusion_metadata_always_present(self, panel_responses, judge_response): + with patch("litellm.acompletion", side_effect=_panel_mock(panel_responses, judge_response)): + result = await afusion(models=PANEL_MODELS, judge_model=JUDGE_MODEL, messages=MESSAGES) + + meta = result._hidden_params["fusion"] + assert meta["judge_model"] == JUDGE_MODEL + assert set(meta["panel_models"]) == set(PANEL_MODELS) + + @pytest.mark.asyncio + async def test_usage_merged(self, panel_responses, judge_response): + with patch("litellm.acompletion", side_effect=_panel_mock(panel_responses, judge_response)): + result = await afusion(models=PANEL_MODELS, judge_model=JUDGE_MODEL, messages=MESSAGES) + + assert result.usage is not None + assert result.usage.total_tokens > 0 + + @pytest.mark.asyncio + async def test_partial_panel_failure_continues(self, panel_responses, judge_response): + async def mock_acompletion(**kwargs): + if kwargs["model"] == "gemini-2.0-flash": + raise RuntimeError("API error") + return judge_response if kwargs["model"] == JUDGE_MODEL else panel_responses[PANEL_MODELS.index(kwargs["model"])] + + with patch("litellm.acompletion", side_effect=mock_acompletion): + result = await afusion(models=PANEL_MODELS, judge_model=JUDGE_MODEL, messages=MESSAGES) + + assert isinstance(result, ModelResponse) + # Only 2 panel models survived + assert len(result._hidden_params["fusion"]["panel_models"]) == 2 + + @pytest.mark.asyncio + async def test_all_panel_fail_raises(self): + async def mock_acompletion(**kwargs): + raise RuntimeError("All failed") + + with patch("litellm.acompletion", side_effect=mock_acompletion): + with pytest.raises(RuntimeError, match="all panel models failed"): + await afusion(models=PANEL_MODELS, judge_model=JUDGE_MODEL, messages=MESSAGES) + + @pytest.mark.asyncio + @pytest.mark.parametrize("strategy", ["single_judge", "majority_vote", "best_of_n"]) + async def test_all_strategies(self, strategy, panel_responses, judge_response): + with patch("litellm.acompletion", side_effect=_panel_mock(panel_responses, judge_response)): + result = await afusion( + models=PANEL_MODELS, + judge_model=JUDGE_MODEL, + messages=MESSAGES, + strategy=strategy, + ) + assert isinstance(result, ModelResponse) + + @pytest.mark.asyncio + async def test_invalid_strategy_raises(self, panel_responses, judge_response): + with patch("litellm.acompletion", side_effect=_panel_mock(panel_responses, judge_response)): + with pytest.raises(ValueError, match="unknown strategy"): + await afusion( + models=PANEL_MODELS, + judge_model=JUDGE_MODEL, + messages=MESSAGES, + strategy="nonexistent", # type: ignore + ) + + +# --------------------------------------------------------------------------- +# Common error cases +# --------------------------------------------------------------------------- + +class TestAFusionKwargs: + """Cover temperature/max_tokens/timeout forwarding and system-prompt preservation.""" + + @pytest.fixture + def panel_responses(self): + return [_make_response(f"Answer {m}", m) for m in PANEL_MODELS] + + @pytest.fixture + def judge_response(self): + return _make_response("Synthesized final answer", JUDGE_MODEL) + + @pytest.mark.asyncio + async def test_temperature_max_tokens_timeout_forwarded(self, panel_responses, judge_response): + """Lines 309, 311, 313, 336: temperature/max_tokens/timeout set on panel_kwargs and judge_kwargs.""" + captured: list[dict] = [] + + async def mock_acompletion(**kwargs): + captured.append(dict(kwargs)) + return judge_response if kwargs["model"] == JUDGE_MODEL else panel_responses[PANEL_MODELS.index(kwargs["model"])] + + with patch("litellm.acompletion", side_effect=mock_acompletion): + await afusion( + models=PANEL_MODELS, + judge_model=JUDGE_MODEL, + messages=MESSAGES, + temperature=0.5, + max_tokens=128, + timeout=10.0, + ) + + panel_calls = [c for c in captured if c["model"] in PANEL_MODELS] + judge_calls = [c for c in captured if c["model"] == JUDGE_MODEL] + assert all(c.get("temperature") == 0.5 for c in panel_calls) + assert all(c.get("max_tokens") == 128 for c in panel_calls) + assert all(c.get("timeout") == 10.0 for c in panel_calls) + assert all(c.get("timeout") == 10.0 for c in judge_calls) + + @pytest.mark.asyncio + async def test_existing_system_prompt_preserved(self, panel_responses, judge_response): + """Lines 126-127: system message in original_messages is carried into judge prompt.""" + messages_with_system = [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "What is 2+2?"}, + ] + captured_judge_messages: list = [] + + async def mock_acompletion(**kwargs): + if kwargs["model"] == JUDGE_MODEL: + captured_judge_messages.extend(kwargs["messages"]) + return judge_response if kwargs["model"] == JUDGE_MODEL else panel_responses[PANEL_MODELS.index(kwargs["model"])] + + with patch("litellm.acompletion", side_effect=mock_acompletion): + await afusion( + models=PANEL_MODELS, + judge_model=JUDGE_MODEL, + messages=messages_with_system, + ) + + system_contents = [m["content"] for m in captured_judge_messages if m.get("role") == "system"] + assert any("helpful assistant" in c for c in system_contents) + + +class TestCommonErrors: + @pytest.mark.asyncio + async def test_empty_models_raises(self): + with pytest.raises(ValueError, match="non-empty list"): + await afusion(models=[], messages=MESSAGES) + + @pytest.mark.asyncio + async def test_empty_models_with_judge_raises(self): + with pytest.raises(ValueError, match="non-empty list"): + await afusion(models=[], judge_model=JUDGE_MODEL, messages=MESSAGES) + + +# --------------------------------------------------------------------------- +# Sync wrapper +# --------------------------------------------------------------------------- + +class TestFusionSync: + def test_callable_via_litellm(self): + assert callable(litellm.fusion) + assert callable(litellm.afusion) + + def test_panel_only_returns_list(self): + panel_resps = [_make_response(f"Panel {m}", m) for m in PANEL_MODELS] + + with patch("litellm.acompletion", side_effect=_panel_mock(panel_resps)): + result = fusion(models=PANEL_MODELS, messages=MESSAGES) + + assert isinstance(result, list) + assert len(result) == len(PANEL_MODELS) + + def test_with_judge_returns_model_response(self): + panel_resps = [_make_response(f"Panel {m}", m) for m in PANEL_MODELS] + judge_resp = _make_response("Sync result", JUDGE_MODEL) + + with patch("litellm.acompletion", side_effect=_panel_mock(panel_resps, judge_resp)): + result = fusion(models=PANEL_MODELS, judge_model=JUDGE_MODEL, messages=MESSAGES) + + assert isinstance(result, ModelResponse) + assert result.choices[0].message.content == "Sync result" + + def test_include_panel_false(self): + panel_resps = [_make_response(f"Panel {m}", m) for m in PANEL_MODELS] + judge_resp = _make_response("Sync result", JUDGE_MODEL) + + with patch("litellm.acompletion", side_effect=_panel_mock(panel_resps, judge_resp)): + result = fusion( + models=PANEL_MODELS, + judge_model=JUDGE_MODEL, + messages=MESSAGES, + include_panel=False, + ) + + assert "panel_responses" not in result._hidden_params["fusion"]