style(fusion): satisfy LIT001/LIT002/LIT010 type-discipline gate

This commit is contained in:
leecoder 2026-08-11 14:55:13 +09:00
parent b1cb97cee4
commit dc72d08974
2 changed files with 115 additions and 99 deletions

View file

@ -1,3 +1,3 @@
from .main import FusionStrategy, afusion, fusion
__all__ = ["FusionStrategy", "afusion", "fusion"]
__all__ = ["FusionStrategy", "afusion", "fusion"] # mutable-ok: module __all__ list, Python convention

View file

@ -22,7 +22,7 @@ Usage:
messages=[{"role": "user", "content": "Explain quantum entanglement"}],
)
# returns: ModelResponse
# response._hidden_params["fusion"]["panel_responses"] ← individual responses
# response._hidden_params["fusion"]["panel_responses"] <- individual responses
# Judge synthesis — panel excluded
response = litellm.fusion(
@ -43,7 +43,8 @@ from __future__ import annotations
import asyncio
import time
import uuid
from typing import Literal, overload
from collections.abc import Sequence
from typing import Final, Literal, TypeAlias, overload
import litellm
from litellm.types.utils import ModelResponse, Usage
@ -52,42 +53,37 @@ from litellm.types.utils import ModelResponse, Usage
# Judge prompt templates
# ---------------------------------------------------------------------------
_JUDGE_PROMPT_SINGLE = """\
You received the following responses from {n} AI models for this user request:
_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"
)
{responses}
_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"
)
Your task: synthesize a single, comprehensive, and accurate final answer.
- Identify the points of agreement and the best reasoning across responses.
- Resolve contradictions by selecting the most well-supported position.
- Fill in gaps where some models provided information others missed.
- Write the final answer as if you generated it directly — no meta-commentary about the synthesis process.
"""
_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_PROMPT_MAJORITY = """\
You received the following responses from {n} AI models for this user request:
{responses}
Select the single best response. Output ONLY the text of the chosen response, verbatim.
"""
_JUDGE_PROMPT_BEST_OF_N = """\
You received the following responses from {n} AI models for this user request:
{responses}
Score each response on a scale of 1-10 for accuracy, completeness, and clarity.
Then output the highest-scored response verbatim.
"""
_JUDGE_PROMPTS: dict[str, str] = {
_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 = Literal["single_judge", "majority_vote", "best_of_n"]
FusionStrategy: TypeAlias = Literal["single_judge", "majority_vote", "best_of_n"]
# ---------------------------------------------------------------------------
@ -96,29 +92,33 @@ FusionStrategy = Literal["single_judge", "majority_vote", "best_of_n"]
def _build_judge_messages(
original_messages: list[dict],
panel_responses: list[ModelResponse],
panel_models: list[str],
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]:
) -> list[dict]: # mutable-ok: message list consumed directly by litellm.acompletion
"""Build the message list for the judge call."""
user_request = next(
(m["content"] for m in reversed(original_messages) if m.get("role") == "user"),
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 = []
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}")
response_blocks.append(f"[Response {i} - {model}]\n{content}")
responses_text = "\n\n".join(response_blocks)
template = _JUDGE_PROMPTS[strategy]
judge_system = template.format(n=len(panel_responses), responses=responses_text)
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] = []
messages: list[dict] = [] # mutable-ok: accumulated then returned immediately # rebind-ok: never rebound
# Preserve any existing system prompt
for m in original_messages:
@ -126,16 +126,16 @@ def _build_judge_messages(
messages.append(m)
break
messages.append({"role": "system", "content": judge_system})
messages.append({"role": "user", "content": user_request})
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: list[ModelResponse]) -> Usage:
def _sum_usage(responses: Sequence[ModelResponse]) -> Usage:
"""Sum token usage across a list of ModelResponse objects."""
total_prompt = 0
total_completion = 0
total_tokens = 0
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
@ -149,25 +149,27 @@ def _sum_usage(responses: list[ModelResponse]) -> Usage:
def _merge_usage(
panel_responses: list[ModelResponse],
panel_responses: Sequence[ModelResponse],
judge_response: ModelResponse,
) -> Usage:
"""Sum token usage across all panel calls + judge call."""
return _sum_usage(panel_responses + [judge_response])
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: list[ModelResponse],
panel_models: list[str],
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 = _merge_usage(panel_responses, judge_response)
merged_usage: Final = _merge_usage(panel_responses, judge_response)
response = ModelResponse(
response: Final = ModelResponse(
id=f"fusion-{uuid.uuid4().hex}",
choices=judge_response.choices,
created=int(time.time()),
@ -176,28 +178,36 @@ def _build_fusion_response(
object="chat.completion",
)
fusion_meta: dict = {
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
fusion_meta["panel_responses"] = panel_responses # mutable-ok: stored in _hidden_params
response._hidden_params = {"fusion": fusion_meta}
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: list[str],
messages: list,
panel_kwargs: dict,
) -> tuple[list[ModelResponse], list[str]]:
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 = [litellm.acompletion(model=m, messages=messages, stream=False, **panel_kwargs) for m in models]
raw: list = list(await asyncio.gather(*tasks, return_exceptions=True))
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] = []
valid_models: list[str] = []
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}")
@ -218,8 +228,8 @@ async def _run_panel(
@overload
async def afusion(
models: list[str],
messages: list,
models: Sequence[str], # mutable-ok: OpenAI-compatible, matches litellm API
messages: Sequence[dict], # mutable-ok: OpenAI-compatible message list
*,
judge_model: None = ...,
strategy: FusionStrategy = ...,
@ -227,14 +237,14 @@ async def afusion(
timeout: float | None = ...,
temperature: float | None = ...,
max_tokens: int | None = ...,
**kwargs: object,
) -> list[ModelResponse]: ...
**kwargs: object, # kwargs-ok: forwarded to litellm.acompletion, varies per provider
) -> list[ModelResponse]: ... # mutable-ok: panel responses list
@overload
async def afusion(
models: list[str],
messages: list,
models: Sequence[str], # mutable-ok: OpenAI-compatible, matches litellm API
messages: Sequence[dict], # mutable-ok: OpenAI-compatible message list
*,
judge_model: str,
strategy: FusionStrategy = ...,
@ -242,7 +252,7 @@ async def afusion(
timeout: float | None = ...,
temperature: float | None = ...,
max_tokens: int | None = ...,
**kwargs: object,
**kwargs: object, # kwargs-ok: forwarded to litellm.acompletion, varies per provider
) -> ModelResponse: ...
@ -252,8 +262,8 @@ async def afusion(
async def afusion(
models: list[str],
messages: list,
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",
@ -261,13 +271,13 @@ async def afusion(
timeout: float | None = None,
temperature: float | None = None,
max_tokens: int | None = None,
**kwargs,
) -> list[ModelResponse] | ModelResponse:
**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).
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.
@ -290,9 +300,11 @@ async def afusion(
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)}")
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)
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:
@ -301,30 +313,34 @@ async def afusion(
panel_kwargs["timeout"] = timeout
# 1. Fan out to all panel models in parallel
valid_responses, valid_models = await _run_panel(models, messages, panel_kwargs)
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 = _build_judge_messages(
judge_messages: Final = _build_judge_messages(
original_messages=messages,
panel_responses=valid_responses,
panel_models=valid_models,
strategy=strategy,
)
judge_kwargs: dict = {}
judge_kwargs: dict = {} # mutable-ok: conditionally extended before use # rebind-ok: never rebound
if timeout is not None:
judge_kwargs["timeout"] = timeout
judge_response: ModelResponse = await litellm.acompletion(
judge_response: Final = await litellm.acompletion(
model=judge_model,
messages=judge_messages,
stream=False,
temperature=0, # deterministic synthesis
**judge_kwargs,
**judge_kwargs, # kwargs-ok: forwarded to litellm.acompletion
)
# 4. Wrap into a single ModelResponse with merged metadata
@ -345,8 +361,8 @@ async def afusion(
@overload
def fusion(
models: list[str],
messages: list,
models: Sequence[str], # mutable-ok: OpenAI-compatible, matches litellm API
messages: Sequence[dict], # mutable-ok: OpenAI-compatible message list
*,
judge_model: None = ...,
strategy: FusionStrategy = ...,
@ -354,14 +370,14 @@ def fusion(
timeout: float | None = ...,
temperature: float | None = ...,
max_tokens: int | None = ...,
**kwargs: object,
) -> list[ModelResponse]: ...
**kwargs: object, # kwargs-ok: forwarded to litellm.acompletion, varies per provider
) -> list[ModelResponse]: ... # mutable-ok: panel responses list
@overload
def fusion(
models: list[str],
messages: list,
models: Sequence[str], # mutable-ok: OpenAI-compatible, matches litellm API
messages: Sequence[dict], # mutable-ok: OpenAI-compatible message list
*,
judge_model: str,
strategy: FusionStrategy = ...,
@ -369,13 +385,13 @@ def fusion(
timeout: float | None = ...,
temperature: float | None = ...,
max_tokens: int | None = ...,
**kwargs: object,
**kwargs: object, # kwargs-ok: forwarded to litellm.acompletion, varies per provider
) -> ModelResponse: ...
def fusion(
models: list[str],
messages: list,
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",
@ -383,8 +399,8 @@ def fusion(
timeout: float | None = None,
temperature: float | None = None,
max_tokens: int | None = None,
**kwargs,
) -> list[ModelResponse] | ModelResponse:
**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.
@ -392,7 +408,7 @@ def fusion(
:func:`afusion` directly.
Args:
models: List of panel model identifiers (≥ 2 recommended).
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.
@ -420,9 +436,9 @@ def fusion(
timeout=timeout,
temperature=temperature,
max_tokens=max_tokens,
**kwargs,
**kwargs, # kwargs-ok: forwarded to afusion
)
)
__all__ = ["FusionStrategy", "afusion", "fusion"]
__all__ = ["FusionStrategy", "afusion", "fusion"] # mutable-ok: module __all__, Python convention