使用到的python库

这里主要用到了这个几个库，使用streamlit，用来页面展示

• streamlit
• langchain
• langchain-core
• requests

具体内容

我们要继承BaseChatModel，并重写_llm_type、_identifying_params、_generate、_stream，以及自定义一个 _format_messages，其中_format_messages是把 LangChain 的消息对象转为 OpenAI 兼容的 {“role”,“content”} 列表。

python代码

CustomChatModel.py

import json
import requests
import httpx
from typing import Any, Dict, List, Iterator, Optional, AsyncIterator, Sequence, Union, Type, Callable

from langchain_core.callbacks import CallbackManagerForLLMRun, AsyncCallbackManagerForLLMRun
from langchain_core.language_models.chat_models import BaseChatModel
from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage, AIMessageChunk, BaseMessage
from langchain_core.messages.tool import ToolCallChunk
from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult
from langchain_core.tools import BaseTool
from langchain_core.utils.function_calling import convert_to_openai_tool
from langchain_core.runnables import Runnable


class CustomChatModel(BaseChatModel):
    base_url: str
    model_name: str
    api_key: str = "sk-3ead9422ee71407586d9bac62251fe76"
    temperature: float = 0.6
    max_tokens: int = 4096  # 稍微调小一点默认值，避免有些模型报错
    streaming: bool = True
    timeout: int = 600

    @property
    def _llm_type(self) -> str:
        return "custom-chat-model"

    @property
    def _identifying_params(self) -> Dict[str, Any]:
        return {
            "model_name": self.model_name,
            "base_url": self.base_url,
            "temperature": self.temperature,
            "max_tokens": self.max_tokens,
        }

    # 实现 bind_tools
    def bind_tools(
            self,
            tools: Sequence[Union[Dict[str, Any], Type, Callable, BaseTool]],
            **kwargs: Any,
    ) -> Runnable[Any, BaseMessage]:
        """将工具绑定到模型（Agent 必须）"""
        formatted_tools = [convert_to_openai_tool(tool) for tool in tools]
        return self.bind(tools=formatted_tools, **kwargs)

    def _format_messages(self, messages: List[BaseMessage]) -> List[Dict[str, Any]]:
        formatted = []
        for msg in messages:
            if isinstance(msg, SystemMessage):
                formatted.append({"role": "system", "content": msg.content})
            elif isinstance(msg, HumanMessage):
                formatted.append({"role": "user", "content": msg.content})
            elif isinstance(msg, AIMessage):
                # 处理 AI 消息，可能包含 content 也可能包含 tool_calls
                msg_dict = {"role": "assistant", "content": msg.content}
                if msg.tool_calls:
                    # 将 LangChain 的 tool_calls 转回 OpenAI 格式
                    tool_calls_payload = []
                    for tc in msg.tool_calls:
                        tool_calls_payload.append({
                            "id": tc["id"],
                            "type": "function",
                            "function": {
                                "name": tc["name"],
                                "arguments": json.dumps(tc["args"])
                            }
                        })
                    msg_dict["tool_calls"] = tool_calls_payload
                formatted.append(msg_dict)
            elif msg.type == "tool":
                # 处理工具执行结果消息 (ToolMessage)
                formatted.append({
                    "role": "tool",
                    "tool_call_id": msg.tool_call_id,
                    "content": msg.content
                })
        return formatted

    def _generate(
            self,
            messages: List[BaseMessage],
            stop: Optional[List[str]] = None,
            run_manager: Optional[CallbackManagerForLLMRun] = None,
            **kwargs: Any,
    ) -> ChatResult:
        # 复用 _stream 的逻辑来生成完整结果
        full_chunk = None
        for chunk in self._stream(messages, stop, run_manager, **kwargs):
            if full_chunk is None:
                full_chunk = chunk
            else:
                full_chunk += chunk

        if not full_chunk:
            return ChatResult(generations=[ChatGeneration(message=AIMessage(content=""))])

        return ChatResult(
            generations=[
                ChatGeneration(
                    message=full_chunk.message
                )
            ]
        )

    def _stream(
            self,
            messages: List[BaseMessage],
            stop: Optional[List[str]] = None,
            run_manager: Optional[CallbackManagerForLLMRun] = None,
            **kwargs: Any,
    ) -> Iterator[ChatGenerationChunk]:
        url = f"{self.base_url}/chat/completions"
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.api_key}"
        }

        # 构建 Payload 时加入 tools
        payload = {
            "model": self.model_name,
            "messages": self._format_messages(messages),
            "temperature": kwargs.get("temperature", self.temperature),
            "max_tokens": kwargs.get("max_tokens", self.max_tokens),
            "stream": True,
        }
        if "tools" in kwargs:
            payload["tools"] = kwargs["tools"]
        if "tool_choice" in kwargs:
            payload["tool_choice"] = kwargs["tool_choice"]
        # ========================================

        try:
            response = requests.post(
                url, headers=headers, json=payload, timeout=self.timeout, stream=True
            )
            response.raise_for_status()

            for line in response.iter_lines():
                if not line: continue
                decoded_line = line.decode("utf-8")
                if not decoded_line.startswith("data:"): continue

                content = decoded_line[len("data: "):].strip()
                if content == "[DONE]": break

                try:
                    chunk_data = json.loads(content)
                    choice = chunk_data.get("choices", [{}])[0]
                    delta = choice.get("delta", {})

                    # 1. 处理思考过程 (Reasoning)
                    reasoning_content = delta.get("reasoning_content")
                    if reasoning_content:
                        chunk = ChatGenerationChunk(
                            message=AIMessageChunk(content="",
                                                   additional_kwargs={"reasoning_content": reasoning_content})
                        )
                        yield chunk
                        if run_manager: run_manager.on_llm_new_token(reasoning_content, chunk=chunk)

                    # 2. 处理正式内容 (Content)
                    content_chunk = delta.get("content")
                    if content_chunk:
                        chunk = ChatGenerationChunk(message=AIMessageChunk(content=content_chunk))
                        yield chunk
                        if run_manager: run_manager.on_llm_new_token(content_chunk, chunk=chunk)

                    # 3. 处理工具调用 (Tool Calls) 
                    tool_calls = delta.get("tool_calls")
                    if tool_calls:
                        tc_chunks = []
                        for tc in tool_calls:
                            tc_chunks.append(ToolCallChunk(
                                name=tc.get("function", {}).get("name"),
                                args=tc.get("function", {}).get("arguments"),
                                id=tc.get("id"),
                                index=tc.get("index")
                            ))
                        chunk = ChatGenerationChunk(
                            message=AIMessageChunk(content="", tool_call_chunks=tc_chunks)
                        )
                        yield chunk

                except (json.JSONDecodeError, KeyError, IndexError):
                    continue

        except requests.RequestException as e:
            raise IOError(f"API request failed: {e}") from e

    async def _astream(
            self,
            messages: List[BaseMessage],
            stop: Optional[List[str]] = None,
            run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
            **kwargs: Any,
    ) -> AsyncIterator[ChatGenerationChunk]:
        url = f"{self.base_url}/chat/completions"
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.api_key}"
        }

        payload = {
            "model": self.model_name,
            "messages": self._format_messages(messages),
            "temperature": kwargs.get("temperature", self.temperature),
            "max_tokens": kwargs.get("max_tokens", self.max_tokens),
            "stream": True,
        }
        if "tools" in kwargs:
            payload["tools"] = kwargs["tools"]
        if "tool_choice" in kwargs:
            payload["tool_choice"] = kwargs["tool_choice"]

        async with httpx.AsyncClient(timeout=self.timeout) as client:
            try:
                async with client.stream("POST", url, headers=headers, json=payload) as response:

                    if response.status_code != 200:
                        await response.aread()  # 注意这里是 aread
                        error_details = response.text
                        raise IOError(f"API HTTP Error {response.status_code}: {error_details}")

                    async for line in response.aiter_lines():
                        if not line: continue
                        if not line.startswith("data:"): continue

                        content = line[len("data: "):].strip()
                        if content == "[DONE]": break

                        try:
                            chunk_data = json.loads(content)
                            choice = chunk_data.get("choices", [{}])[0]
                            delta = choice.get("delta", {})

                            # 1. 处理思考过程 (Reasoning)
                            reasoning_content = delta.get("reasoning_content")
                            if reasoning_content:
                                chunk = ChatGenerationChunk(
                                    message=AIMessageChunk(content="",
                                                           additional_kwargs={"reasoning_content": reasoning_content})
                                )
                                yield chunk
                                if run_manager: await run_manager.on_llm_new_token(reasoning_content, chunk=chunk)

                            # 2. 处理正式内容 (Content)
                            content_chunk = delta.get("content")
                            if content_chunk:
                                chunk = ChatGenerationChunk(message=AIMessageChunk(content=content_chunk))
                                yield chunk
                                if run_manager: await run_manager.on_llm_new_token(content_chunk, chunk=chunk)

                            # 3. 处理工具调用
                            tool_calls = delta.get("tool_calls")
                            if tool_calls:
                                tc_chunks = []
                                for tc in tool_calls:
                                    tc_chunks.append(ToolCallChunk(
                                        name=tc.get("function", {}).get("name"),
                                        args=tc.get("function", {}).get("arguments"),
                                        id=tc.get("id"),
                                        index=tc.get("index")
                                    ))
                                chunk = ChatGenerationChunk(
                                    message=AIMessageChunk(content="", tool_call_chunks=tc_chunks)
                                )
                                yield chunk

                        except (json.JSONDecodeError, KeyError, IndexError):
                            continue

            except httpx.HTTPStatusError as e:
                await e.response.aread()  # 同样使用 aread
                raise IOError(f"HTTP Status Error: {e.response.text}")
            except httpx.RequestError as e:
                raise IOError(f"Async API request failed: {e}")

app.py

import streamlit as st
from langchain_core.messages import HumanMessage, SystemMessage, AIMessage
from typing import List, Dict, Any

# 导入您提供的自定义模型
try:
    from CustomChatModel import CustomChatModel
except ImportError:
    st.error("错误：未找到 `CustomChatModel.py` 文件。请确保它与 app.py 位于同一目录。")
    st.stop()

# --- 配置模型 ---
BASE_URL = "http://192.168.1.89:8075/v1" 
MODEL_NAME = "Qwen3-235B-A22B-Thinking-Pro6000"
CHAT_SYSTEM_PROMPT = """你是一个智能对话助手。
- 你的回答必须是 Markdown 格式。"""

# -----------------

@st.cache_resource
def load_llm():
    """
    缓存 LLM 实例，避免每次都重新加载。
    """
    try:
        return CustomChatModel(
            base_url=BASE_URL,
            model_name=MODEL_NAME
        )
    except Exception as e:
        st.error(f"初始化模型失败: {e}")
        return None

def format_messages(chat_history: List[Dict[str, str]]) -> List:
    """
    将 Streamlit 聊天记录转换为 LangChain 消息格式。
    """
    # 注意：在这个 Demo 中，我们没有文档上下文，所以系统提示很简单
    # 简化的系统提示
    simple_system_prompt = "你是一个乐于助人的 AI 助手。"
    
    messages = [SystemMessage(content=simple_system_prompt)]
    
    for msg in chat_history:
        if msg["role"] == "user":
            messages.append(HumanMessage(content=msg["content"]))
        else:
            messages.append(AIMessage(content=msg["content"]))
    return messages


# --- Streamlit 界面 ---

st.title("🤖 简易聊天 Demo")
st.caption(f"正在连接到: `{MODEL_NAME}` at `{BASE_URL}`")

# 加载模型
llm = load_llm()
if llm is None:
    st.stop()

# 初始化聊天记录
if "messages" not in st.session_state:
    st.session_state.messages = []

# 显示历史消息
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

# 接收用户输入
if prompt := st.chat_input("请输入您的问题..."):
    # 将用户消息添加到聊天记录
    st.session_state.messages.append({"role": "user", "content": prompt})
    
    # 显示用户消息
    with st.chat_message("user"):
        st.markdown(prompt)

    # --- AI 响应逻辑 ---
    with st.chat_message("assistant"):
        # 创建一个空占位符，用于流式更新
        placeholder = st.empty()
        
        full_response = ""
        thinking_content = ""
        
        # 1. 准备 LangChain 消息
        formatted_messages = format_messages(st.session_state.messages)
        
        # 2. 从 views.py 的 chat 视图中获取模型参数
        model_kwargs = {
            "max_tokens": 8192,
            "temperature": 0.3,
        }

        try:
            # 3. 调用模型的 .stream() 方法
            stream = llm.stream(formatted_messages, **model_kwargs)

            # 4. 迭代处理流式响应块
            for chunk in stream:
                # 检查模型特有的 "thinking" 内容
                reasoning_chunk = chunk.additional_kwargs.get("reasoning_content")
                if reasoning_chunk:
                    thinking_content += reasoning_chunk
                    # 在占位符中显示 "thinking" 状态
                    placeholder.markdown(f"🤔 *{thinking_content}*")
                
                # 检查标准 "content" 内容
                content_chunk = chunk.content
                if content_chunk:
                    # 如果这是 "thinking" 后的第一个内容块，清空 "thinking" 消息
                    if thinking_content and full_response == "":
                        placeholder.empty()
                    
                    full_response += content_chunk
                    # 实时更新占位符中的内容
                    placeholder.markdown(full_response + "▌") # "▌" 是一个模拟光标
            
            # 5. 流式结束后，显示完整内容（不带光标）
            placeholder.markdown(full_response)

        except Exception as e:
            st.error(f"调用模型时出错: {e}")

    # 将 AI 的完整响应添加到聊天记录
    if full_response:
        st.session_state.messages.append({"role": "assistant", "content": full_response})