RAG实战6-如何在LlamaIndex使用自己搭建的大模型API

在搭建一个大模型API服务中，我们介绍了如何使用SWIFT框架搭建一个大模型API服务。在RAG实战1-5中，我们一直使用的是本地加载大模型的方式来调用大模型，本文将介绍如何在LlamaIndex中使用自己搭建的大模型API。

LlamaIndex支持部分厂商的API配置，如OpenAI，但我们想使用的是自己在服务器上搭建的API服务，这个时候需要我们定制一个LLM类，代码如下：

from typing import Any
from llama_index.core import PromptTemplate, Settings, StorageContext, load_index_from_storage
from llama_index.core.base.llms.types import LLMMetadata, CompletionResponse, CompletionResponseGen
from llama_index.core.llms import CustomLLM
from llama_index.core.llms.callbacks import llm_completion_callback
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from swift.llm import get_model_list_client, XRequestConfig, inference_client

# API
model_list = get_model_list_client()
model_type = model_list.data[0].id
print(f'API model_type: {model_type}')
request_config = XRequestConfig(seed=42)


# 定制自己的LLM类
class BianCangLLM(CustomLLM):
    context_window: int = 4096
    num_output: int = 2048
    model_name: str = "BianCang"

    @property
    def metadata(self) -> LLMMetadata:
        """Get LLM metadata."""
        return LLMMetadata(
            context_window=self.context_window,
            num_output=self.num_output,
            model_name=self.model_name,
        )

    @llm_completion_callback()
    def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
        resp = inference_client(model_type, prompt, [], request_config=request_config)
        return CompletionResponse(text=resp.choices[0].message.content)

    @llm_completion_callback()
    def stream_complete(
        self, prompt: str, **kwargs: Any
    ) -> CompletionResponseGen:
        resp = inference_client(model_type, prompt, [], request_config=request_config)
        response = ""
        for token in resp.choices[0].message.content:
            response += token
            yield CompletionResponse(text=response, delta=token)


# 定义system prompt
SYSTEM_PROMPT = """你是一个医疗人工智能助手。"""
query_wrapper_prompt = PromptTemplate(
    "[INST]<<SYS>>\n" + SYSTEM_PROMPT + "<</SYS>>\n\n{query_str}[/INST] "
)

# 定义qa prompt
qa_prompt_tmpl_str = (
    "上下文信息如下。\n"
    "---------------------\n"
    "{context_str}\n"
    "---------------------\n"
    "请根据上下文信息而不是先验知识来回答以下的查询。"
    "作为一个医疗人工智能助手，你的回答要尽可能严谨。\n"
    "Query: {query_str}\n"
    "Answer: "
)
qa_prompt_tmpl = PromptTemplate(qa_prompt_tmpl_str)

# 使用自定义的LLM API
Settings.llm = BianCangLLM()

# 使用llama-index-embeddings-huggingface构建本地embedding模型
Settings.embed_model = HuggingFaceEmbedding(
    model_name="E:\\LLMs\\bge-base-zh-v1.5"
)

# 从存储文件中读取embedding向量和向量索引
storage_context = StorageContext.from_defaults(persist_dir="doc_emb")
index = load_index_from_storage(storage_context)

# 构建查询引擎
query_engine = index.as_query_engine(similarity_top_k=5)

# 更新查询引擎中的prompt template
query_engine.update_prompts(
    {"response_synthesizer:text_qa_template": qa_prompt_tmpl}
)

# 查询获得答案
response = query_engine.query("不耐疲劳，口燥、咽干可能是哪些证候？")
print(response)

代码的核心是实现BianCangLLM类，该类继承自LlamaIndex的CustomLLM类。我们需要重写父类中的def metadata(self) -> LLMMetadata、def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse、def stream_complete(self, prompt: str, **kwargs: Any) -> CompletionResponseGen:。其中，metadata负责定义大模型的一些参数属性；complete负责调用大模型API服务并直接返回响应；stream_complete负责调用大模型API服务并以流式输出的形式返回响应。

运行代码，同样可以得到类似于之前的效果：

根据提供的上下文信息，口燥、咽干的症状可能与以下几个中医证候相关：

1. 津液不足证（4.6.1.1）：由于津液生成不足或体内燥热，可能导致口眼喉鼻干燥。

2. 津亏热结证（4.6.3.2）：津液亏虚加上热邪内结，也会出现口燥咽干的表现。

3. 津液亏涸证（4.6.1.2）：津液亏损严重时，口干、唇裂、鼻燥、舌燥是其特征，可能包括咽干。

4. 燥干清窍证（3.6.3.2）：气候干燥导致的津液耗损，会引起口鼻咽喉干燥。

5. 津伤化燥证（6.3.1）：燥热内蕴或内热化燥可能引起口干舌燥，伴有多尿、消瘦等症状。

因此，这些证候都可能与不耐疲劳和口燥咽干的临床表现相关，但具体诊断需要结合其他症状和中医辨证原则。建议患者就诊中医师以获取专业诊断。

好处是，我们不需要每次启动RAG应用时都加载一遍大模型权重了。