251 lines
9.3 KiB
Python
251 lines
9.3 KiB
Python
"""LiteLLM 封装: capabilities 决定调用参数,自动重试。
|
||
|
||
`chat()`:同步阻塞,一次性返回完整 response。给 probe / 离线探测用。
|
||
`chat_stream()`:流式 generator,yield chunk;调用方累积 + 用 litellm.stream_chunk_builder
|
||
拼回完整 response。loop 传 cancel_check 后由独立 pump 承担 provider 阻塞读取，
|
||
控制线程固定节拍响应停止，不依赖下一块到达。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import os
|
||
import queue
|
||
import threading
|
||
import time
|
||
from typing import Any, Callable, Iterator, List, Optional
|
||
|
||
# 跳过启动时从 GitHub 拉 model_prices 的网络请求,直接用 litellm 打包的本地副本。
|
||
# 必须在 `import litellm` 之前设置,否则 get_model_cost_map() 已经跑过了。
|
||
os.environ.setdefault("LITELLM_LOCAL_MODEL_COST_MAP", "True")
|
||
|
||
import litellm # noqa: E402
|
||
from litellm.exceptions import (
|
||
APIConnectionError,
|
||
APIError,
|
||
RateLimitError,
|
||
ServiceUnavailableError,
|
||
Timeout,
|
||
)
|
||
|
||
from .capabilities import ModelCapabilities
|
||
from .llm_params import build_thinking_kwargs
|
||
|
||
# 单次 LLM 请求超时(秒),默认与 litellm 一致(600s)但显式化 + env 可调 ──
|
||
# 长思考模型真被掐("600s 无字节 → run 标 error")时调大 ZCBOT_LLM_TIMEOUT_S 即可,
|
||
# 不用改代码。流式场景它主要约束"无字节间隔",正常出 chunk 不触发。
|
||
# 长工具调用不受此限:>1min 的脚本走 background=true(DESIGN §8.12),不占 LLM 超时。
|
||
_REQUEST_TIMEOUT_S = int(os.getenv("ZCBOT_LLM_TIMEOUT_S", "600"))
|
||
|
||
|
||
class LLM:
|
||
def __init__(self, capabilities: ModelCapabilities) -> None:
|
||
self.caps = capabilities
|
||
env_name = capabilities.api_key_env or "DEEPSEEK_API_KEY"
|
||
self.api_key = os.environ.get(env_name)
|
||
self.api_base = capabilities.api_base or None
|
||
if not self.api_key:
|
||
raise RuntimeError(
|
||
f"环境变量 {env_name} 未设置,无法调用 {capabilities.model_id}"
|
||
)
|
||
|
||
def _build_kwargs(
|
||
self,
|
||
messages: List[dict],
|
||
tools: Optional[list],
|
||
parallel_tool_calls: Optional[bool],
|
||
reasoning_effort: Optional[str],
|
||
) -> dict:
|
||
kwargs: dict = {
|
||
"model": self.caps.model_id,
|
||
"messages": messages,
|
||
"temperature": self.caps.optimal_temperature,
|
||
"api_key": self.api_key,
|
||
"timeout": _REQUEST_TIMEOUT_S,
|
||
}
|
||
if self.api_base:
|
||
kwargs["api_base"] = self.api_base
|
||
if tools:
|
||
kwargs["tools"] = tools
|
||
if self.caps.parallel_tools and parallel_tool_calls is not False:
|
||
kwargs["parallel_tool_calls"] = True
|
||
kwargs.update(
|
||
build_thinking_kwargs(
|
||
enabled=self.caps.thinking_enabled,
|
||
transport=self.caps.thinking_transport,
|
||
reasoning_effort=reasoning_effort,
|
||
)
|
||
)
|
||
if self.caps.prompt_caching:
|
||
kwargs["extra_headers"] = {"anthropic-beta": "prompt-caching-2024-07-31"}
|
||
return kwargs
|
||
|
||
def chat(
|
||
self,
|
||
messages: List[dict],
|
||
tools: Optional[list] = None,
|
||
parallel_tool_calls: Optional[bool] = None,
|
||
reasoning_effort: Optional[str] = None,
|
||
max_retries: int = 3,
|
||
) -> Any:
|
||
kwargs = self._build_kwargs(messages, tools, parallel_tool_calls, reasoning_effort)
|
||
last_err: Optional[Exception] = None
|
||
for attempt in range(max_retries):
|
||
try:
|
||
response = litellm.completion(**kwargs)
|
||
return response
|
||
except (RateLimitError, APIConnectionError, ServiceUnavailableError, Timeout, APIError) as e:
|
||
last_err = e
|
||
if attempt == max_retries - 1:
|
||
break
|
||
time.sleep(2 ** attempt)
|
||
raise last_err # type: ignore[misc]
|
||
|
||
def chat_stream(
|
||
self,
|
||
messages: List[dict],
|
||
tools: Optional[list] = None,
|
||
parallel_tool_calls: Optional[bool] = None,
|
||
reasoning_effort: Optional[str] = None,
|
||
max_retries: int = 3,
|
||
cancel_check: Optional[Callable[[], bool]] = None,
|
||
) -> Iterator[Any]:
|
||
"""流式 chat:yield 每个 chunk。调用方累积 + 用 litellm.stream_chunk_builder 拼回完整 response。
|
||
|
||
重试语义:连接建立阶段错误(还没拿到第一个 chunk)按 max_retries 退避重试;
|
||
开始流之后失败直接抛(半截 partial 没法续)。usage 通过 stream_options.include_usage
|
||
让最后一个 chunk 带 usage。
|
||
|
||
传入 cancel_check 时，底层阻塞读取在 daemon pump 线程执行，当前 run 线程按
|
||
100ms 轮询取消。这样 provider 在 TTFT / reasoning 分片之间长时间无字节时，
|
||
停止也不必等到下一块到达；已有底层 stream 会被主动 close。
|
||
"""
|
||
kwargs = self._build_kwargs(messages, tools, parallel_tool_calls, reasoning_effort)
|
||
kwargs["stream"] = True
|
||
kwargs["stream_options"] = {"include_usage": True}
|
||
|
||
if cancel_check is not None:
|
||
yield from self._chat_stream_interruptible(
|
||
kwargs, max_retries=max_retries, cancel_check=cancel_check,
|
||
)
|
||
return
|
||
|
||
yield from self._chat_stream_direct(kwargs, max_retries=max_retries)
|
||
|
||
@staticmethod
|
||
def _chat_stream_direct(kwargs: dict, *, max_retries: int) -> Iterator[Any]:
|
||
"""无取消调用方的直连路径（probe 等保持原有同步语义）。"""
|
||
|
||
last_err: Optional[Exception] = None
|
||
for attempt in range(max_retries):
|
||
try:
|
||
stream = litellm.completion(**kwargs)
|
||
break
|
||
except (RateLimitError, APIConnectionError, ServiceUnavailableError, Timeout, APIError) as e:
|
||
last_err = e
|
||
if attempt == max_retries - 1:
|
||
raise
|
||
time.sleep(2 ** attempt)
|
||
else:
|
||
raise last_err # type: ignore[misc]
|
||
|
||
try:
|
||
for chunk in stream:
|
||
yield chunk
|
||
finally:
|
||
# 调用方提前 break(cancel) → generator close → 这里关掉底层 httpx 连接
|
||
close = getattr(stream, "close", None)
|
||
if callable(close):
|
||
try:
|
||
close()
|
||
except Exception:
|
||
pass
|
||
|
||
@staticmethod
|
||
def _chat_stream_interruptible(
|
||
kwargs: dict,
|
||
*,
|
||
max_retries: int,
|
||
cancel_check: Callable[[], bool],
|
||
) -> Iterator[Any]:
|
||
"""把 provider 的阻塞迭代与 run 控制线程解耦，以固定节拍响应取消。"""
|
||
items: queue.Queue[tuple[str, Any]] = queue.Queue(maxsize=64)
|
||
stop = threading.Event()
|
||
stream_box: dict[str, Any] = {}
|
||
|
||
def offer(kind: str, value: Any = None) -> bool:
|
||
while not stop.is_set():
|
||
try:
|
||
items.put((kind, value), timeout=0.1)
|
||
return True
|
||
except queue.Full:
|
||
continue
|
||
return False
|
||
|
||
def pump() -> None:
|
||
stream = None
|
||
try:
|
||
last_err: Optional[Exception] = None
|
||
for attempt in range(max_retries):
|
||
if stop.is_set():
|
||
return
|
||
try:
|
||
stream = litellm.completion(**kwargs)
|
||
stream_box["stream"] = stream
|
||
break
|
||
except (
|
||
RateLimitError,
|
||
APIConnectionError,
|
||
ServiceUnavailableError,
|
||
Timeout,
|
||
APIError,
|
||
) as e:
|
||
last_err = e
|
||
if attempt == max_retries - 1:
|
||
raise
|
||
if stop.wait(2 ** attempt):
|
||
return
|
||
else:
|
||
if last_err is not None:
|
||
raise last_err
|
||
return
|
||
|
||
for chunk in stream:
|
||
if not offer("chunk", chunk):
|
||
return
|
||
offer("done")
|
||
except BaseException as e: # noqa: BLE001 - 原样转抛到 run 线程
|
||
offer("error", e)
|
||
finally:
|
||
if stream is not None:
|
||
close = getattr(stream, "close", None)
|
||
if callable(close):
|
||
try:
|
||
close()
|
||
except Exception:
|
||
pass
|
||
|
||
worker = threading.Thread(target=pump, daemon=True, name="llm-stream-pump")
|
||
worker.start()
|
||
try:
|
||
while True:
|
||
if cancel_check():
|
||
return
|
||
try:
|
||
kind, value = items.get(timeout=0.1)
|
||
except queue.Empty:
|
||
continue
|
||
if kind == "chunk":
|
||
yield value
|
||
elif kind == "error":
|
||
raise value
|
||
else:
|
||
return
|
||
finally:
|
||
stop.set()
|
||
stream = stream_box.get("stream")
|
||
close = getattr(stream, "close", None)
|
||
if callable(close):
|
||
try:
|
||
close()
|
||
except Exception:
|
||
pass
|