|
|
| from __future__ import annotations
|
|
|
| import os
|
| from dataclasses import dataclass, field
|
| from typing import Any
|
|
|
| from huggingface_hub import InferenceClient
|
| from langchain_core.messages import AIMessage
|
| from langchain_core.runnables import RunnableConfig
|
| from langchain_openai import ChatOpenAI
|
|
|
| from ..completion import complete_turn
|
| from ..config import INFERENCE_MODE, MODEL_ID, local_inference_enabled
|
| from ..messages import api_turn_to_ai_message, langchain_messages_to_api
|
|
|
| HF_ROUTER_BASE_URL = "https://router.huggingface.co/v1"
|
|
|
|
|
| def _resolve_hf_token(configurable: dict[str, Any]) -> str:
|
| oauth_token = str(configurable.get("hf_token") or "").strip()
|
| if oauth_token:
|
| return oauth_token
|
| return os.environ.get("HF_TOKEN", "").strip()
|
|
|
|
|
| @dataclass
|
| class HubInferenceChatModel:
|
| """LangGraph-compatible wrapper around huggingface_hub InferenceClient."""
|
|
|
| hf_token: str
|
| max_tokens: int = 1800
|
| temperature: float = 0.35
|
| top_p: float = 0.9
|
| tools: list[dict[str, Any]] = field(default_factory=list)
|
|
|
| def bind_tools(self, tools: list[dict[str, Any]]) -> HubInferenceChatModel:
|
| return HubInferenceChatModel(
|
| hf_token=self.hf_token,
|
| max_tokens=self.max_tokens,
|
| temperature=self.temperature,
|
| top_p=self.top_p,
|
| tools=list(tools),
|
| )
|
|
|
| def invoke(self, messages: list[Any]) -> AIMessage:
|
| client = InferenceClient(api_key=self.hf_token, model=MODEL_ID)
|
| api_messages = langchain_messages_to_api(messages)
|
| content, reasoning, tool_calls = complete_turn(
|
| client,
|
| api_messages,
|
| max_tokens=self.max_tokens,
|
| temperature=self.temperature,
|
| top_p=self.top_p,
|
| tools=self.tools or None,
|
| )
|
| return api_turn_to_ai_message(content, reasoning, tool_calls)
|
|
|
|
|
| @dataclass
|
| class MiniCPMChatModel:
|
| max_tokens: int = 1800
|
| temperature: float = 0.35
|
| top_p: float = 0.9
|
| tools: list[dict[str, Any]] = field(default_factory=list)
|
|
|
| def bind_tools(self, tools: list[dict[str, Any]]) -> MiniCPMChatModel:
|
| return MiniCPMChatModel(
|
| max_tokens=self.max_tokens,
|
| temperature=self.temperature,
|
| top_p=self.top_p,
|
| tools=list(tools),
|
| )
|
|
|
| def invoke(self, messages: list[Any]) -> AIMessage:
|
| from ..minicpm.model import chat_complete
|
|
|
| return chat_complete(
|
| messages,
|
| tools=self.tools or None,
|
| max_tokens=self.max_tokens,
|
| temperature=self.temperature,
|
| top_p=self.top_p,
|
| )
|
|
|
|
|
| def build_llm(
|
| config: RunnableConfig,
|
| **overrides: Any,
|
| ) -> HubInferenceChatModel | MiniCPMChatModel | ChatOpenAI:
|
| """Build a chat model from the per-request configurable values."""
|
| configurable = config.get("configurable", {})
|
| max_tokens = int(overrides.pop("max_tokens", configurable.get("max_tokens", 1800)))
|
| temperature = float(
|
| overrides.pop("temperature", configurable.get("temperature", 0.35))
|
| )
|
| top_p = float(overrides.pop("top_p", configurable.get("top_p", 0.9)))
|
| hf_token = _resolve_hf_token(configurable)
|
|
|
| if local_inference_enabled():
|
| return MiniCPMChatModel(
|
| max_tokens=max_tokens,
|
| temperature=temperature,
|
| top_p=top_p,
|
| )
|
|
|
| if INFERENCE_MODE == "router":
|
| params: dict[str, Any] = {
|
| "model": MODEL_ID,
|
| "api_key": hf_token,
|
| "base_url": HF_ROUTER_BASE_URL,
|
| "max_tokens": max_tokens,
|
| "temperature": temperature,
|
| "top_p": top_p,
|
| }
|
| params.update(overrides)
|
| return ChatOpenAI(**params)
|
|
|
| return HubInferenceChatModel(
|
| hf_token=hf_token,
|
| max_tokens=max_tokens,
|
| temperature=temperature,
|
| top_p=top_p,
|
| )
|
|
|
