LangGraph入门(01:状态)

LangChain入门到实战(七:RAG)

Milvus作为向量数据库中的热门选型之一，本篇文章带你了解下Milvus


从业务角度，Milvus数据模型层级如下
`Database -> Collection -> Partition -> Entity`
- Database：Milvus 的数据库，用来隔离不同业务数据。
- Collection：最核心的逻辑容器，类似于关系型数据库里的 table。
- Partition：分区，是 collection 的子集，不是必须手动创建；一个 collection 至少会有默认partition。
- Entity：可以理解为 collection 中的一条记录。类似于关系型数据库的一行数据。


# 1. DDL操作

## 1.1 Database
创建客户端
```python
from pymilvus import MilvusClient
client = MilvusClient("http://localhost:19530")
```
列出所有数据库
```python
existed_databases = client.list_databases()

for db in existed_databases:
    print(db)
```

创建数据库
```python
db_name = "rag_demo"

if db_name not in existed_databases:
    client.create_database(db_name=db_name)
```

删除数据库

如果数据库下有Collection则无法删除，需要先删除它的所有Collection才能删除Database
```python
client.drop_database(db_name=db_name)
```
切换数据库
```python
client.use_database(db_name=db_name)
```

## 1.2 Collection
查看数据库下的collections
```python

collections = client.list_collections()

for coll in collections:
    print(coll)
```

创建collection
```python
collection_name = "docs"

client.create_collection(
    collection_name=collection_name,
    dimension=1024,
    metric_type="COSINE"
)
```
- dimension ：嵌入向量维度，应和嵌入模型的向量维度保持一致，见下文。
- metric_type 表示向量相似度的计算方式， COSINE 表示余弦相似度。

    当用户提问时，系统会把提问也变成向量，然后去数据库里找“最相似”的本地文本向量。但怎么定
    义“相似”呢？ 向量数据库需要知道计算规则。 COSINE （余弦相似度）关注的是两个向量在方向上的夹角：
    - 如果两个向量方向完全一致（代表文本意思极度接近），余弦值接近 1 。
    - 如果方向正交（毫无关系），值接近 0 。
    - 结论：在 RAG 检索时，Milvus 会帮你计算用户问题与数据库中所有文本的余弦相似度，并把
    得分（Score）从大到小排序，把得分最高（最相似）的前 K 个片段还给你。 (注：除了COSINE 之外，常见的还有 L2 欧氏距离、 IP 内积等。)
    
 删除collection
 ```python
 client.drop_collection(collection_name=collection_name)
 ```
 
 # 2. DML操作
 ## 2.1 准备嵌入模型
 ```python
 import os

# HF国内镜像
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
# 防止CPU多线程死锁
os.environ["OMP_NUM_THREADS"] = "1"
os.environ["MKL_NUM_THREADS"] = "1"
from typing import List

import torch
from dotenv import load_dotenv
from langchain_core.embeddings import Embeddings
from pydantic import BaseModel, PrivateAttr
from transformers import AutoTokenizer, AutoModel

load_dotenv(override=True)


class BGESentenceEmbedding(BaseModel, Embeddings):
    model_name: str
    cache_dir: str
    device: str = "cpu"

    # ✅ PrivateAttr：实例私有变量，不参与pydantic校验，完美解决冲突
    _tokenizer: AutoTokenizer = PrivateAttr()
    _model: AutoModel = PrivateAttr()

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        print("开始加载tokenizer...")
        self._tokenizer = AutoTokenizer.from_pretrained(
            self.model_name,
            cache_dir=self.cache_dir
        )
        print("开始加载model...")
        self._model = AutoModel.from_pretrained(
            self.model_name,
            cache_dir=self.cache_dir
        ).to(self.device)
        self._model.eval()
        print("✅ 模型加载完毕！")

    @staticmethod
    def _mean_pooling(model_output, attention_mask):
        token_embeddings = model_output[0]
        input_mask = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
        return torch.sum(token_embeddings * input_mask, 1) / torch.clamp(input_mask.sum(1), min=1e-9)

    def embed_query(self, text: str) -> List[float]:
        encoded = self._tokenizer(
            text, padding=True, truncation=True, return_tensors="pt"
        ).to(self.device)
        with torch.no_grad():
            model_output = self._model(**encoded)
        embeddings = self._mean_pooling(model_output, encoded["attention_mask"])
        embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)
        return embeddings[0].cpu().tolist()

    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        return [self.embed_query(t) for t in texts]


cache_path = os.path.join(os.getcwd(), "../embeddings")
embedding_model = BGESentenceEmbedding(
    model_name="BAAI/bge-base-zh-v1.5",
    cache_dir=cache_path,
    device="cpu"
)
 ```
 ## 2.2 准备collection
 创建collection
 ```python
 collection_name = "docs"

client.create_collection(
    collection_name=collection_name,
    dimension=768, 
    metric_type="COSINE"
)
```
`BAAI/bge-base-zh-v1.5`的维度是768,要和嵌入模型的保持一致，不然在数据存入时会报错


 查看collection元数据,
 ```python
 
 from rich import print as rprint

metadata = client.describe_collection(collection_name=collection_name)

rprint(metadata)
```


![image.png](/static/img/e8af14cf6323ccddd98db3f6ac70430e.image.webp)
我们在创建collection时没有指定schema，后者可以理解为表结构，此时Milvus会将collection定义为默认结构。字段信息如下
- id ：数据ID，作为主键唯一标识数据
- vector ：数据的嵌入向量

此外， `enable_dynamic_field` 为 `True` ，这表示 `collection` 支持动态字段。也就是说，除了预定义的`id` 和 `vector` 字段之外，在插入数据时还可以携带其他未提前声明的字段，这些字段会被自动写入并 统一存储在动态字段 中。这样做的好处是能够在不修改 `schema` 的情况下，灵活保存额外的业务属性，例如文本内容、标签、时间戳或来源信息等，适合字段结构不固定的场景。

## 2.3 准备数据
```python
# 准备测试数据
texts = [
    "LangChain 是一个用于构建 LLM 应用的开发框架。",
    "Milvus 是一个适合 AI 应用的向量数据库。",
    "RAG 的核心是先检索相关知识，再让大模型生成答案。",
    "Docker Desktop 可以方便地在本地运行 Milvus Standalone。"
]

# 生成嵌入向量
vectors = embedding_model.embed_documents(texts)

# 封装可以存到milvus的数据
data = [
    {
        "id": i,
        "vector": vectors[i],
        "text": texts[i],
        "source": "demo"
    } for i in range(len(texts))
]
```

## 2.4 写入数据
```python

insert_res = client.upsert(
    collection_name=collection_name,
    data=data,
)

print("insert result : ", insert_res)
# insert result :  {'upsert_count': 4, 'ids': [0, 1, 2, 3]}
```
upsert可以保证幂等写入，即主键相同时覆盖

手动flush
```python
client.flush(collection_name=collection_name)
```
Milvus不会第一时间将数据落盘，要看到写入效果，我们手动flush，将数据刷写到磁盘

查看collection统计信息
```python
stats = client.get_collection_stats(collection_name=collection_name)

print("stats : ", stats)
```

# 3. DQL操作

## 3.1 扫描数据
```python
iterator = client.query_iterator(
    collection_name=collection_name,
    filter="",
    output_fields=["*"]
)

i = 0
while True:

    rows = iterator.next()

    if not rows:
        break

    for row in rows:
        print(f"第{i + 1}条数据：\n")

        print(f'id : {row["id"]},vector = {row["vector"][:5]},text = {row["text"]},source = {row["source"]}')

        #i += 1

iterator.close()
```

## 3.2 通过主键查询数据
```python
res = client.get(
    collection_name=collection_name,
    ids=[0, 1, 2]
)

print(len(res))

for i in range(len(res)):
    print(f"第{i + 1}条数据：")
    print(f'id : {res[i]["id"]},vector = {res[i]["vector"][:5]},text = {res[i]["text"]},source = {res[i]["source"]}')
    # print(res[i])

```

## 3.3 相似度检索
```python
# 相似度检索
query = "什么是向量数据库？"
query_vector = embedding_model.embed_query(query)

results = client.search(
    collection_name=collection_name,
    data=[query_vector],
    limit=3,
    output_fields=["text", "source", "id"]
)

for res in results[0]:
    print(res)
```
检索结果如下：
```python
{'id': 1, 'distance': 0.6177977919578552, 'entity': {'id': 1, 'text': 'Milvus 是一个适合 AI 应用的向量数据库。', 'source': 'demo'}}
{'id': 0, 'distance': 0.46440669894218445, 'entity': {'id': 0, 'text': 'LangChain 是一个用于构建 LLM 应用的开发框架。', 'source': 'demo'}}
{'id': 3, 'distance': 0.3925101161003113, 'entity': {'id': 3, 'text': 'Docker Desktop 可以方便地在本地运行 Milvus Standalone。', 'source': 'demo'
```