491 lines
19 KiB
Python
491 lines
19 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
from copy import deepcopy
|
|
from dataclasses import dataclass, field
|
|
from typing import Any, Dict, Iterable, Iterator, List, Optional, Tuple
|
|
from langchain_core.language_models.chat_models import BaseChatModel
|
|
from .utils import ensure_env_loaded
|
|
try:
|
|
# pydantic v2
|
|
from pydantic import PrivateAttr
|
|
except Exception: # pragma: no cover
|
|
# pydantic v1 fallback
|
|
from pydantic.fields import PrivateAttr # type: ignore
|
|
|
|
from langchain_core.utils.function_calling import convert_to_openai_tool
|
|
|
|
|
|
def _env(name: str, default: Optional[str] = None) -> Optional[str]:
|
|
v = os.getenv(name)
|
|
return v if v else default
|
|
|
|
|
|
def _coalesce(*vals):
|
|
for v in vals:
|
|
if v is not None:
|
|
return v
|
|
return None
|
|
|
|
|
|
def _to_openai_messages(messages: List["BaseMessage"]) -> List[Dict[str, Any]]:
|
|
from langchain_core.messages import (
|
|
AIMessage,
|
|
BaseMessage,
|
|
HumanMessage,
|
|
SystemMessage,
|
|
ToolMessage,
|
|
FunctionMessage,
|
|
)
|
|
|
|
out: List[Dict[str, Any]] = []
|
|
for m in messages:
|
|
if isinstance(m, SystemMessage):
|
|
out.append({"role": "system", "content": m.content})
|
|
elif isinstance(m, HumanMessage):
|
|
out.append({"role": "user", "content": m.content})
|
|
elif isinstance(m, AIMessage):
|
|
# If previous AI produced tool calls, pass them through
|
|
entry: Dict[str, Any] = {"role": "assistant", "content": m.content}
|
|
tcs = getattr(m, "tool_calls", None) or m.additional_kwargs.get("tool_calls") if hasattr(m, "additional_kwargs") else None
|
|
if tcs:
|
|
# Map to OpenAI's expected shape
|
|
entry["tool_calls"] = [
|
|
{
|
|
"id": tc.get("id") if isinstance(tc, dict) else getattr(tc, "id", None),
|
|
"type": "function",
|
|
"function": {
|
|
"name": (tc.get("name") if isinstance(tc, dict) else getattr(tc, "name", "")) or "",
|
|
"arguments": json.dumps(tc.get("args") if isinstance(tc, dict) else getattr(tc, "args", {})),
|
|
},
|
|
}
|
|
for tc in (tcs or [])
|
|
]
|
|
out.append(entry)
|
|
elif isinstance(m, ToolMessage):
|
|
# OpenAI-compatible "tool" role with name + id must be provided
|
|
out.append(
|
|
{
|
|
"role": "tool",
|
|
"content": m.content,
|
|
"tool_call_id": m.tool_call_id,
|
|
}
|
|
)
|
|
elif isinstance(m, FunctionMessage):
|
|
# Legacy function messages
|
|
out.append(
|
|
{
|
|
"role": "tool",
|
|
"content": m.content,
|
|
"tool_call_id": m.name or "",
|
|
}
|
|
)
|
|
else:
|
|
# Fallback
|
|
out.append({"role": getattr(m, "role", "user"), "content": getattr(m, "content", str(m))})
|
|
return out
|
|
|
|
|
|
def _normalize_params(schema: Any) -> Dict[str, Any]:
|
|
"""Normalize various schema-like inputs into a valid OpenAI JSON Schema."""
|
|
if not isinstance(schema, dict):
|
|
return {"type": "object", "properties": {}, "required": []}
|
|
if schema.get("type") == "object" and isinstance(schema.get("properties"), dict):
|
|
out = dict(schema)
|
|
elif "properties" in schema and isinstance(schema["properties"], dict):
|
|
out = {"type": "object", **schema}
|
|
elif all(isinstance(v, dict) for v in schema.values()):
|
|
out = {"type": "object", "properties": schema}
|
|
else:
|
|
out = {"type": "object", "properties": {}}
|
|
for k in ("$schema", "$defs", "definitions", "title"):
|
|
out.pop(k, None)
|
|
req = out.get("required", [])
|
|
if not isinstance(req, list):
|
|
req = list(req) if req else []
|
|
out["required"] = req
|
|
return out
|
|
|
|
|
|
def _convert_tool(tool: Any) -> Dict[str, Any]:
|
|
"""Convert LangChain/MCP tool to OpenAI JSON via official utility, with fallback normalization."""
|
|
try:
|
|
t = convert_to_openai_tool(tool)
|
|
fn = t.get("function", {})
|
|
fn["description"] = (fn.get("description") or "")[:1000]
|
|
fn["parameters"] = _normalize_params(fn.get("parameters"))
|
|
return {"type": "function", "function": fn}
|
|
except Exception:
|
|
pass
|
|
name = getattr(tool, "name", None) or getattr(tool, "__name__", "tool")
|
|
desc = getattr(tool, "description", "") or ""
|
|
schema = None
|
|
if getattr(tool, "args_schema", None) is not None:
|
|
try:
|
|
schema = tool.args_schema.model_json_schema()
|
|
except Exception:
|
|
try:
|
|
schema = tool.args_schema.schema()
|
|
except Exception:
|
|
schema = None
|
|
if schema is None and hasattr(tool, "args"):
|
|
try:
|
|
schema = tool.args
|
|
except Exception:
|
|
schema = None
|
|
params = _normalize_params(schema)
|
|
return {"type": "function", "function": {"name": name, "description": desc[:1000], "parameters": params}}
|
|
|
|
|
|
def _convert_tools(tools: Optional[List[Any]]) -> Optional[List[Dict[str, Any]]]:
|
|
if not tools:
|
|
return None
|
|
return [_convert_tool(t) for t in tools]
|
|
|
|
|
|
def _parse_tool_calls(msg: Dict[str, Any]) -> Tuple[str, List[Dict[str, Any]]]:
|
|
content = _coalesce(msg.get("content"), "")
|
|
tool_calls_raw = msg.get("tool_calls") or []
|
|
parsed: List[Dict[str, Any]] = []
|
|
for tc in tool_calls_raw:
|
|
fn = (tc or {}).get("function", {})
|
|
name = fn.get("name") or ""
|
|
args_str = fn.get("arguments") or "{}"
|
|
args = {}
|
|
if isinstance(args_str, str):
|
|
try:
|
|
args = json.loads(args_str or "{}")
|
|
except Exception:
|
|
# keep raw if malformed; downstream may repair
|
|
args = {"$raw": args_str}
|
|
elif isinstance(args_str, dict):
|
|
args = args_str
|
|
parsed.append({
|
|
"id": tc.get("id"),
|
|
"name": name,
|
|
"args": args,
|
|
})
|
|
return content, parsed
|
|
|
|
|
|
@dataclass
|
|
class _ToolCallBuilder:
|
|
id: Optional[str] = None
|
|
name: str = ""
|
|
args_buf: List[str] = field(default_factory=list)
|
|
|
|
def to_final(self) -> Dict[str, Any]:
|
|
args_str = "".join(self.args_buf) if self.args_buf else "{}"
|
|
try:
|
|
args = json.loads(args_str or "{}")
|
|
except Exception:
|
|
args = {"$raw": args_str}
|
|
return {"id": self.id, "name": self.name, "args": args}
|
|
|
|
|
|
class ChatQwenOpenAICompat(BaseChatModel):
|
|
"""
|
|
A minimal LangChain BaseChatModel-compatible class that calls an
|
|
OpenAI-compatible /chat/completions endpoint and assembles tool calls.
|
|
|
|
Notes:
|
|
- Implements _generate for non-stream and _stream for SSE streaming.
|
|
- Stores bound tools and tool_choice to send with requests.
|
|
- Avoids heavy dependencies; uses httpx at runtime.
|
|
"""
|
|
|
|
# Private state to avoid pydantic field validation
|
|
_model: str = PrivateAttr()
|
|
_api_key: str = PrivateAttr()
|
|
_base_url: str = PrivateAttr()
|
|
_temperature: Optional[float] = PrivateAttr(default=None)
|
|
_max_tokens: Optional[int] = PrivateAttr(default=None)
|
|
_extra_body: Dict[str, Any] = PrivateAttr(default_factory=dict)
|
|
_timeout: float = PrivateAttr(default=60.0)
|
|
_tools: Optional[List[Dict[str, Any]]] = PrivateAttr(default=None)
|
|
_tool_choice: Optional[str] = PrivateAttr(default=None)
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
model: Optional[str] = None,
|
|
api_key: Optional[str] = None,
|
|
base_url: Optional[str] = None,
|
|
temperature: Optional[float] = None,
|
|
max_tokens: Optional[int] = None,
|
|
extra_body: Optional[Dict[str, Any]] = None,
|
|
timeout: float = 60.0,
|
|
) -> None:
|
|
super().__init__()
|
|
# Load .env variables once if available
|
|
ensure_env_loaded()
|
|
# Store config in private attrs
|
|
self._model = model or _env("QWEN_MODEL", "qwen3-coder-flash")
|
|
self._api_key = (
|
|
api_key
|
|
or _env("QWEN_API_KEY")
|
|
or _env("GPUSTACK_API_KEY")
|
|
or _env("OPENAI_API_KEY")
|
|
or _env("DASHSCOPE_API_KEY")
|
|
)
|
|
if not self._api_key:
|
|
raise ValueError("Missing API key: set QWEN_API_KEY/OPENAI_API_KEY/DASHSCOPE_API_KEY")
|
|
self._base_url = base_url or _env("QWEN_BASE_URL") or _env("OPENAI_BASE_URL")
|
|
if not self._base_url:
|
|
# Default to DashScope intl OpenAI-compatible
|
|
self._base_url = "https://dashscope-intl.aliyuncs.com/compatible-mode/v1"
|
|
# Allow env override for timeout
|
|
timeout_env = _env("QWEN_TIMEOUT")
|
|
if timeout_env:
|
|
try:
|
|
timeout = float(timeout_env)
|
|
except Exception:
|
|
pass
|
|
self._temperature = temperature
|
|
self._max_tokens = max_tokens
|
|
self._extra_body = extra_body or {}
|
|
self._timeout = timeout
|
|
# tools/tool_choice already defaulted via PrivateAttr declarations
|
|
|
|
# Expose LangChain binding-style API
|
|
def bind_tools(self, tools: List[Any]) -> "ChatQwenOpenAICompat":
|
|
clone = self._copy()
|
|
clone._tools = _convert_tools(tools)
|
|
return clone
|
|
|
|
def bind(self, **kwargs: Any) -> "ChatQwenOpenAICompat":
|
|
clone = self._copy()
|
|
# Allow passing tool_choice or extra_body overrides
|
|
if "tool_choice" in kwargs:
|
|
clone._tool_choice = kwargs["tool_choice"]
|
|
if "extra_body" in kwargs and isinstance(kwargs["extra_body"], dict):
|
|
eb = dict(clone._extra_body)
|
|
eb.update(kwargs["extra_body"]) # type: ignore
|
|
clone._extra_body = eb
|
|
return clone
|
|
|
|
# LangChain BaseChatModel required interface
|
|
@property
|
|
def _llm_type(self) -> str: # type: ignore[override]
|
|
return "qwen-openai-compat"
|
|
|
|
def _copy(self) -> "ChatQwenOpenAICompat":
|
|
c = ChatQwenOpenAICompat(
|
|
model=self._model,
|
|
api_key=self._api_key,
|
|
base_url=self._base_url,
|
|
temperature=self._temperature,
|
|
max_tokens=self._max_tokens,
|
|
extra_body=dict(self._extra_body),
|
|
timeout=self._timeout,
|
|
)
|
|
c._tools = self._tools[:] if self._tools else None
|
|
c._tool_choice = self._tool_choice
|
|
return c
|
|
|
|
# -- Internal HTTP helpers --
|
|
def _request(self, payload: Dict[str, Any], *, stream: bool = False) -> Any:
|
|
try:
|
|
import httpx
|
|
import sys as _sys
|
|
except Exception as e:
|
|
raise RuntimeError("httpx is required for ChatQwenOpenAICompat. Install via `pip install httpx`. ") from e
|
|
|
|
# Auth header customization (for non-standard gateways)
|
|
auth_header = _env("QWEN_AUTH_HEADER", "Authorization")
|
|
auth_scheme = _env("QWEN_AUTH_SCHEME", "Bearer")
|
|
auth_value = f"{auth_scheme} {self._api_key}" if auth_scheme else self._api_key
|
|
headers = {auth_header: auth_value, "Content-Type": "application/json"}
|
|
|
|
base = self._base_url.rstrip("/")
|
|
url = base if "chat/completions" in base else base + "/chat/completions"
|
|
|
|
# honor system proxy by default; disable via QWEN_HTTP_TRUST_ENV=0
|
|
trust_env = _env("QWEN_HTTP_TRUST_ENV", "1") != "0"
|
|
|
|
if _env("QWEN_DEBUG") == "1":
|
|
print(f"[qwen] POST {url} timeout={self._timeout} trust_env={trust_env}", file=_sys.stderr)
|
|
|
|
# ===== DEBUG: dump payload & headers (API key redacted) =====
|
|
if _env("QWEN_DEBUG_BODY") == "1":
|
|
try:
|
|
# 打印请求体
|
|
print("[qwen] payload =\n" + json.dumps(payload, ensure_ascii=False, indent=2), file=_sys.stderr)
|
|
except Exception:
|
|
print("[qwen] payload = <non-serializable>", file=_sys.stderr)
|
|
try:
|
|
# 打印打码后的 headers
|
|
redacted = dict(headers)
|
|
if auth_header in redacted and isinstance(redacted[auth_header], str):
|
|
# 只保留认证方案，密钥打码
|
|
val = redacted[auth_header]
|
|
if " " in val:
|
|
scheme, _token = val.split(" ", 1)
|
|
redacted[auth_header] = f"{scheme} ***REDACTED***"
|
|
else:
|
|
redacted[auth_header] = "***REDACTED***"
|
|
print("[qwen] headers = " + json.dumps(redacted, ensure_ascii=False), file=_sys.stderr)
|
|
except Exception:
|
|
pass
|
|
# ===========================================================
|
|
|
|
client = httpx.Client(timeout=self._timeout, trust_env=trust_env)
|
|
|
|
if stream:
|
|
# 流式时也会先把 payload 打出来（上面已做），这里直接返回流
|
|
return client.stream("POST", url, headers=headers, json=payload)
|
|
|
|
resp = client.post(url, headers=headers, json=payload)
|
|
|
|
# ===== DEBUG: dump response =====
|
|
if _env("QWEN_DEBUG_RESP") == "1":
|
|
print(f"[qwen] status={resp.status_code}", file=_sys.stderr)
|
|
try:
|
|
print("[qwen] resp.text =\n" + resp.text, file=_sys.stderr)
|
|
except Exception:
|
|
print("[qwen] resp.text = <binary or non-text>", file=_sys.stderr)
|
|
# =================================
|
|
|
|
try:
|
|
resp.raise_for_status()
|
|
except httpx.HTTPStatusError as e:
|
|
# 打印错误体，帮助定位 4xx/5xx 的具体原因
|
|
if _env("QWEN_DEBUG_RESP", "1") == "1":
|
|
try:
|
|
print(f"[qwen] ERROR status={e.response.status_code}\n{e.response.text}", file=_sys.stderr)
|
|
except Exception:
|
|
pass
|
|
raise
|
|
return resp.json()
|
|
|
|
|
|
def _build_payload(self, messages: List["BaseMessage"], **kwargs: Any) -> Dict[str, Any]:
|
|
payload: Dict[str, Any] = {
|
|
"model": self._model,
|
|
"messages": _to_openai_messages(messages),
|
|
}
|
|
if self._temperature is not None:
|
|
payload["temperature"] = self._temperature
|
|
if self._max_tokens is not None:
|
|
payload["max_tokens"] = self._max_tokens
|
|
if self._tools:
|
|
safe_tools: List[Dict[str, Any]] = []
|
|
for t in self._tools:
|
|
if isinstance(t, dict) and t.get("type") == "function" and isinstance(t.get("function"), dict):
|
|
fn = deepcopy(t["function"]) # type: ignore
|
|
fn["description"] = (fn.get("description") or "")[:1000]
|
|
fn["parameters"] = _normalize_params(fn.get("parameters"))
|
|
safe_tools.append({"type": "function", "function": fn})
|
|
else:
|
|
safe_tools.append(_convert_tool(t))
|
|
payload["tools"] = safe_tools
|
|
if self._tool_choice is not None:
|
|
payload["tool_choice"] = self._tool_choice
|
|
if self._extra_body:
|
|
payload.update(self._extra_body)
|
|
if kwargs:
|
|
allowed = {
|
|
"tool_choice","temperature","top_p","max_tokens","stop","stream","n","logprobs","logit_bias","user","presence_penalty","frequency_penalty","seed","parallel_tool_calls","response_format"
|
|
}
|
|
for k, v in kwargs.items():
|
|
if k in allowed and v is not None:
|
|
payload[k] = v
|
|
return payload
|
|
|
|
# -- Non-streaming generate --
|
|
def _generate(
|
|
self,
|
|
messages: List["BaseMessage"],
|
|
stop: Optional[List[str]] = None,
|
|
run_manager: Optional[Any] = None,
|
|
**kwargs: Any,
|
|
) -> "ChatResult": # type: ignore[name-defined]
|
|
from langchain_core.outputs import ChatGeneration, ChatResult
|
|
from langchain_core.messages import AIMessage
|
|
|
|
payload = self._build_payload(messages, **kwargs)
|
|
data = self._request(payload, stream=False)
|
|
choice = (data.get("choices") or [{}])[0]
|
|
msg = choice.get("message", {})
|
|
content, tool_calls = _parse_tool_calls(msg)
|
|
usage = data.get("usage", {})
|
|
|
|
extra_kwargs = {
|
|
"raw": data,
|
|
"reasoning_content": msg.get("reasoning_content"),
|
|
}
|
|
msg_kwargs: Dict[str, Any] = {"content": content or "", "additional_kwargs": extra_kwargs}
|
|
if tool_calls:
|
|
msg_kwargs["tool_calls"] = tool_calls
|
|
ai = AIMessage(**msg_kwargs)
|
|
gen = ChatGeneration(message=ai)
|
|
return ChatResult(generations=[gen], llm_output={"usage": usage, "model": data.get("model")})
|
|
|
|
# -- Streaming generate --
|
|
def _stream(
|
|
self,
|
|
messages: List["BaseMessage"],
|
|
stop: Optional[List[str]] = None,
|
|
run_manager: Optional[Any] = None,
|
|
**kwargs: Any,
|
|
) -> Iterator["ChatGenerationChunk"]: # type: ignore[name-defined]
|
|
from langchain_core.outputs import ChatGenerationChunk
|
|
from langchain_core.messages import AIMessageChunk
|
|
|
|
payload = self._build_payload(messages, stream=True, **kwargs)
|
|
with self._request(payload, stream=True) as r: # type: ignore[attr-defined]
|
|
# SSE stream; lines like: data: {...}\n\n; end with data: [DONE]
|
|
text_buf: List[str] = []
|
|
tc_builders: List[_ToolCallBuilder] = []
|
|
for line in r.iter_lines(): # type: ignore[attr-defined]
|
|
if not line:
|
|
continue
|
|
if isinstance(line, (bytes, bytearray)):
|
|
try:
|
|
line = line.decode("utf-8", errors="ignore")
|
|
except Exception:
|
|
continue
|
|
if not line.startswith("data:"):
|
|
continue
|
|
data = line[len("data:"):].strip()
|
|
if data == "[DONE]":
|
|
# Finalize: emit a chunk with final tool calls if any
|
|
if tc_builders or text_buf:
|
|
extra: Dict[str, Any] = {}
|
|
if tc_builders:
|
|
extra["tool_calls"] = [b.to_final() for b in tc_builders]
|
|
final = AIMessageChunk(
|
|
content="".join(text_buf) if text_buf else "",
|
|
additional_kwargs=extra if extra else None,
|
|
)
|
|
yield ChatGenerationChunk(message=final)
|
|
break
|
|
try:
|
|
obj = json.loads(data)
|
|
except Exception:
|
|
continue
|
|
choices = obj.get("choices") or []
|
|
if not choices:
|
|
continue
|
|
delta = (choices[0].get("delta") or {})
|
|
# Text deltas
|
|
dtext = delta.get("content")
|
|
if dtext:
|
|
text_buf.append(dtext)
|
|
yield ChatGenerationChunk(message=AIMessageChunk(content=dtext))
|
|
# Tool call deltas
|
|
dtools = delta.get("tool_calls") or []
|
|
for i, part in enumerate(dtools):
|
|
# Ensure list length
|
|
while len(tc_builders) <= i:
|
|
tc_builders.append(_ToolCallBuilder())
|
|
b = tc_builders[i]
|
|
if "id" in part:
|
|
b.id = part.get("id") or b.id
|
|
fn = part.get("function") or {}
|
|
if fn.get("name"):
|
|
b.name = fn["name"]
|
|
if fn.get("arguments"):
|
|
b.args_buf.append(fn["arguments"]) # partial JSON string
|