9行代码开发一个基于ollama的私有化RAG

前言

6行代码入门RAG开发 提到其中的OpenAI（LLM + Embedding）是使用LiteLLM + ollama模拟，具体做法如下，

Llamaindex OpenAI LLM 模型
默认使用的是gpt-3.5-turbo
， embedding 模型
默认使用的是text-embedding-ada-002
， 所以这里使用litellm 配合config来模拟，

一句话介绍LiteLLM

使用OpenAI格式调用所有LLM API。使用Bedrock，Azure，OpenAI，Cohere，Anthropic，Ollama，Sagemaker，HuggingFace，Replicate（100+ LLM）。

安装LiteLLM

pip install 'litellm[proxy]'

模型映射

创建litellm.yaml

model_list: 
  - model_name:  gpt-3.5-turbo 
    litellm_params: 
      model: ollama/wizardlm2:7b-q5_K_M
      max_tokens: 8192
  - model_name: text-embedding-ada-002
    litellm_params: 
      model: ollama/snowflake-arctic-embed:latest
      api_base: http://localhost:11434
litellm_settings: # module level litellm settings 
  drop_params: True

运行litellm

litellm --config litellm.yaml

💡，这样模拟的OpenAI不支持function call，所以index.as_chat_engine(chat_mode="condense_question")
 不能使用默认的chat_mode。

如果想部署并使用开源模型，不需要这么麻烦，我们可以根据需要选择合适的方案，这里我介绍下ollama的实现。

Ollama

Ollama 提供了多种LLM模型和embedding模型（如all-minilm
,mxbai-embed-large
,nomic-embed-text
,snowflake-arctic-embed
 等)， 如果没有你想用的，也可以自己导入gguf格式模型。

关于ollama，这里就不展开了，有兴趣的可以看看我的另外一篇文章：ollama 使用技巧集锦

安装包

pip install llama-index-embeddings-ollama
pip install llama-index-llms-ollama

代码

from llama_index.core import VectorStoreIndex, Document, SimpleDirectoryReader,Settings
from llama_index.llms.ollama import Ollama
from llama_index.embeddings.ollama import OllamaEmbedding
# 指定LLM
Settings.llm = Ollama(model="wizardlm2:7b-q5_K_M", request_timeout=60.0)
# 指定 embedding model
Settings.embed_model = OllamaEmbedding(model_name="snowflake-arctic-embed:latest")
## 剩下代码一样
documents = SimpleDirectoryReader("./data").load_data()
index = VectorStoreIndex.from_documents(documents)
chat_engine = index.as_chat_engine(chat_mode="condense_question", verbose=True)
print(chat_engine.chat("DuckDB的VSS扩展主要功能, reply in Chinese"))

结论

至此，一个基于ollama的rag就有雏形了。