mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-08 03:08:45 +00:00
style(fusion): satisfy LIT001/LIT002/LIT010 type-discipline gate
This commit is contained in:
parent
b1cb97cee4
commit
dc72d08974
2 changed files with 115 additions and 99 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue