使用 Milvus 与 Llama Stack 构建 RAG

Llama Stack 是一种面向服务、API 优先的方法，用于构建生产 AI 应用程序。它提供了一个通用堆栈，允许开发人员在任何地方开发、在任何地方部署，并利用具有真正提供商独立性的生产就绪构建块。Llama Stack 专注于 Meta 的 Llama 模型、可组合性、生产就绪性和合作伙伴生态系统。

在本教程中，我们将介绍如何构建配置了 Milvus 的 Llama Stack 服务器，使您能够导入私有数据作为知识库。然后我们将在服务器上执行查询，创建一个完整的 RAG 应用程序。

准备环境

启动 Llama Stack 服务器有多种方式，例如作为库、构建分发包等。对于 Llama Stack 中的每个组件，也可以选择各种提供商。因此，有很多种方式来启动 Llama Stack 服务器。

本教程使用以下配置作为示例来启动服务。如果您希望以其他方式启动，请参考启动 Llama Stack 服务器。

• 我们使用 Conda 构建带有 Milvus 配置的自定义分发包。
• 我们使用 Together AI 作为 LLM 提供商。
• 我们使用默认的 all-MiniLM-L6-v2 作为嵌入模型。

本教程主要参考 Llama Stack 文档 的官方安装指南。如果您在本教程中发现任何过时的部分，您可以优先遵循官方指南并为我们创建一个 issue。

启动 Llama Stack 服务器

准备环境

由于我们需要使用 Together AI 作为 LLM 服务，我们必须首先登录官方网站申请一个 API 密钥，并将 API 密钥 TOGETHER_API_KEY 设置为环境变量。

克隆 Llama Stack 源代码

$ git clone https://github.com/meta-llama/llama-stack.git
$ cd llama-stack

创建 conda 环境并安装依赖

$ conda create -n stack python=3.10
$ conda activate stack

$ pip install -e .

修改 llama_stack/llama_stack/template/together/run.yaml 中的内容，将 vector_io 部分更改为相关的 Milvus 配置。例如，添加：

vector_io:
- provider_id: milvus
  provider_type: inline::milvus
  config:
    db_path: ~/.llama/distributions/together/milvus_store.db

#  - provider_id: milvus
#    provider_type: remote::milvus
#    config:
#      uri: http://localhost:19530
#      token: root:Milvus

在 Llama Stack 中，Milvus 可以通过两种方式配置：本地配置，即 inline::milvus，和远程配置，即 remote::milvus。

• 最简单的方法是本地配置，需要设置 db_path，这是本地存储 Milvus-Lite 文件的路径。

• 远程配置适用于大数据存储。

  • 如果您有大量数据，可以在 Docker 或 Kubernetes 上设置高性能的 Milvus 服务器。在此设置中，请使用服务器 URI，例如 http://localhost:19530 作为您的 uri。默认的 token 是 root:Milvus。
  • 如果您想使用 Zilliz Cloud，Milvus 的完全托管云服务，请调整 uri 和 token，它们对应于 Zilliz Cloud 中的公共端点和 API 密钥。

从模板构建分发包

运行以下命令构建分发包：

$ llama stack build --template together --image-type conda

将在 ~/.llama/distributions/together/together-run.yaml 生成一个文件。然后，运行此命令启动服务器：

$ llama stack run --image-type conda ~/.llama/distributions/together/together-run.yaml

如果一切顺利，您应该看到 Llama Stack 服务器成功在端口 8321 上运行。

从客户端执行 RAG

启动服务器后，您可以编写客户端代码来访问它。以下是示例代码：

import uuid
from llama_stack_client.types import Document
from llama_stack_client.lib.agents.agent import Agent
from llama_stack_client.types.agent_create_params import AgentConfig

# See https://www.together.ai/models for all available models
INFERENCE_MODEL = "meta-llama/Llama-3.3-70B-Instruct-Turbo"
LLAMA_STACK_PORT = 8321


def create_http_client():
    from llama_stack_client import LlamaStackClient

    return LlamaStackClient(
        base_url=f"http://localhost:{LLAMA_STACK_PORT}"  # Your Llama Stack Server URL
    )


client = create_http_client()

# Documents to be used for RAG
urls = ["chat.rst", "llama3.rst", "memory_optimizations.rst", "lora_finetune.rst"]
documents = [
    Document(
        document_id=f"num-{i}",
        content=f"https://raw.githubusercontent.com/pytorch/torchtune/main/docs/source/tutorials/{url}",
        mime_type="text/plain",
        metadata={},
    )
    for i, url in enumerate(urls)
]

# Register a vector database
vector_db_id = f"test-vector-db-{uuid.uuid4().hex}"
client.vector_dbs.register(
    vector_db_id=vector_db_id,
    embedding_model="all-MiniLM-L6-v2",
    embedding_dimension=384,
    provider_id="milvus",
)

print("inserting...")
# Insert the documents into the vector database
client.tool_runtime.rag_tool.insert(
    documents=documents, vector_db_id=vector_db_id, chunk_size_in_tokens=1024,
)

agent_config = AgentConfig(
    model=INFERENCE_MODEL,
    # Define instructions for the agent ( aka system prompt)
    instructions="You are a helpful assistant",
    enable_session_persistence=False,
    # Define tools available to the agent
    toolgroups=[{"name": "builtin::rag", "args": {"vector_db_ids": [vector_db_id]}}],
)

rag_agent = Agent(client, agent_config)
session_id = rag_agent.create_session("test-session")
print("finish init agent...")
user_prompt = (
    "What are the top 5 topics that were explained? Only list succinct bullet points."
)

# Get the final answer from the agent
response = rag_agent.create_turn(
    messages=[{"role": "user", "content": user_prompt}],
    session_id=session_id,
    stream=False,
)
print(f"Response: ")
print(response.output_message.content)

运行此代码来执行 RAG 查询。 如果一切正常工作，输出应该如下所示：

inserting...
finish init agent...
Response: 
* Fine-Tuning Llama3 with Chat Data
* Evaluating fine-tuned Llama3-8B models with EleutherAI's Eval Harness
* Generating text with our fine-tuned Llama3 model
* Faster generation via quantization
* Fine-tuning on a custom chat dataset