RAG项目第三集：自定义LLM大模型类的构建

本集我们通过第一集提供的大模型api接口，开发一个基于langchain的LLM的继承类，代码如下

from langchain.llms.base import LLM
from typing import Any, List, Optional
from langchain.callbacks.manager import CallbackManagerForLLMRun
import requests


class CustomLLM(LLM):
    endpoint: str = "***/v1"  # 本地大模型的API地址
    model: str = "***" #大模型名称

    @property
    def _llm_type(self) -> str:
        return "***" #大模型名称

    def _call(
        self,
        prompt: str,
        stop: Optional[List[str]] = None,
        callbacks: Optional[CallbackManagerForLLMRun] = None,
        **kwargs: Any,
        ) -> str:
        headers = {"Content-Type": "application/json"}
        data = {"model": self.model, "messages": [{"role": "user", "content": prompt}]}
        response = requests.post(f"{self.endpoint}/chat/completions", headers=headers, json=data)
        if response.status_code != 200:
            return "error"

        result = response.json()
        text = result["choices"][0]["message"]["content"]
        return text

通过定义这个类，我们可以调用本地大模型了， 为RAG主函数打下了基础

这个项目代码我已经开源到GitHub上了，可以点这个链接，直接下载整个项目代码："github.com/davidhandso…"

接下来就是第四集：RAG功能的实现