如何为LangChain自定义LLM实现bind_tools与_stream方法？

1. 实现bind_tools方法解决报错 #

create_tool_calling_agent要求LLM类必须实现bind_tools方法，用于将工具信息绑定到模型，让模型知晓可用工具及调用格式。以下是修改后的完整实现：

核心修改点 #

• 添加bound_tools属性存储绑定的工具
• 实现bind_tools方法返回绑定工具后的LLM实例
• 在_call和_stream方法中自动将工具描述和调用规则添加到prompt，引导模型生成符合要求的工具调用格式

from langchain.llms.base import LLM
from langchain.tools import BaseTool
from langchain.callbacks.manager import CallbackManagerForLLMRun
from langchain.schema import GenerationChunk
from typing import List, Optional, Any, Iterator, Dict
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

model_id = "meta-llama/Meta-Llama-3-8B-Instruct"

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,
    device_map="auto",
)

# 为tokenizer设置pad token（避免generate报错）
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

def llama3_instruct(prompt, stream=False):
    # 按照Llama3 Instruct官方要求格式化prompt
    formatted_prompt = (
        "<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n"
        "You are a helpful assistant. Follow instructions carefully.<|eot_id|>"
        f"<|start_header_id|>user<|end_header_id|>\n{prompt}<|eot_id|>"
        "<|start_header_id|>assistant<|end_header_id|>\n"
    )
    
    input_ids = tokenizer.encode(formatted_prompt, return_tensors='pt').to(model.device)
    
    terminators = [
        tokenizer.eos_token_id,
        tokenizer.convert_tokens_to_ids("<|eot_id|>")
    ]
    
    if stream:
        # 真正的流式生成，逐token返回结果
        for output in model.generate(
            input_ids,
            max_new_tokens=256,
            eos_token_id=terminators,
            do_sample=True,
            temperature=0.6,
            top_p=0.9,
            stream=True,
            pad_token_id=tokenizer.pad_token_id
        ):
            new_token = output[0][-1:]
            decoded_token = tokenizer.decode(new_token, skip_special_tokens=True)
            if decoded_token:
                yield decoded_token
    else:
        outputs = model.generate(
            input_ids,
            max_new_tokens=256,
            eos_token_id=terminators,
            do_sample=True,
            temperature=0.6,
            top_p=0.9,
            pad_token_id=tokenizer.pad_token_id
        )
        response = outputs[0][input_ids.shape[-1]:]
        return tokenizer.decode(response, skip_special_tokens=True)

class CustomLLM(LLM):
    bound_tools: Optional[List[BaseTool]] = None

    def _call(
        self,
        prompt: str,
        stop: Optional[List[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> str:
        modified_prompt = prompt
        # 如果绑定了工具，将工具描述和调用规则加入prompt
        if self.bound_tools:
            tools_desc = "\n".join([f"- {tool.name}: {tool.description}" for tool in self.bound_tools])
            tool_call_rules = """
当需要调用工具时，请严格按照以下JSON格式输出（不要添加其他无关内容）：
{
  "name": "工具名称",
  "parameters": {
    "参数名": 参数值
  }
}
无需调用工具时直接回答问题即可。
"""
            modified_prompt = f"""你可以使用以下工具辅助回答问题：
{tools_desc}
{tool_call_rules}

{prompt}"""
        
        res = llama3_instruct(modified_prompt)
        
        # 处理stop词，截断结果
        if stop:
            for stop_word in stop:
                if stop_word in res:
                    res = res[:res.index(stop_word)]
        return res

    def _stream(
        self,
        prompt: str,
        stop: Optional[List[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> Iterator[GenerationChunk]:
        modified_prompt = prompt
        if self.bound_tools:
            tools_desc = "\n".join([f"- {tool.name}: {tool.description}" for tool in self.bound_tools])
            tool_call_rules = """
当需要调用工具时，请严格按照以下JSON格式输出（不要添加其他无关内容）：
{
  "name": "工具名称",
  "parameters": {
    "参数名": 参数值
  }
}
无需调用工具时直接回答问题即可。
"""
            modified_prompt = f"""你可以使用以下工具辅助回答问题：
{tools_desc}
{tool_call_rules}

{prompt}"""
        
        current_text = ""
        for token in llama3_instruct(modified_prompt, stream=True):
            current_text += token
            
            # 实时检查stop词，触发则立即停止生成
            if stop:
                for stop_word in stop:
                    if stop_word in current_text:
                        split_idx = current_text.index(stop_word)
                        if split_idx > 0:
                            chunk = GenerationChunk(text=current_text[:split_idx])
                            if run_manager:
                                run_manager.on_llm_new_token(chunk.text, chunk=chunk)
                            yield chunk
                        return
            
            # 返回当前token的生成块
            chunk = GenerationChunk(text=token)
            if run_manager:
                run_manager.on_llm_new_token(chunk.text, chunk=chunk)
            yield chunk

    def bind_tools(self, tools: List[BaseTool], **kwargs: Any) -> "CustomLLM":
        # 返回绑定工具后的新LLM实例，避免修改原实例
        return CustomLLM(bound_tools=tools)

    @property
    def _identifying_params(self) -> Dict[str, Any]:
        return {
            "model_name": "Meta-Llama-3-8B-Instruct",
            "bound_tools": [tool.name for tool in self.bound_tools] if self.bound_tools else None
        }

    @property
    def _llm_type(self) -> str:
        return "custom_llama3"

2. 优化_stream方法说明 #

原_stream方法是先调用_call获取完整结果再逐字符拆分，并非真正的流式生成。优化后的实现：

• 直接调用模型的流式生成接口（stream=True），逐token返回结果，降低延迟
• 实时检查stop词，触发时立即停止生成，符合LangChain的流式规范
• 完全贴合Llama3的输出逻辑，避免不必要的内存占用

3. 其他注意事项 #

• 必须为tokenizer设置pad_token，否则模型生成时会抛出错误
• 严格按照Llama3 Instruct的官方prompt格式构造输入，能大幅提升模型响应的准确性
• create_tool_calling_agent会自动调用bind_tools方法，无需手动调用

内容的提问来源于stack exchange，提问作者AVJ