第6章: 私有知识库搭建：文档切分、向量化与入库

# Python 3.x (推荐)
pip3 install dashscope

# 或者使用 pip
pip install dashscope

import dashscope
from dashscope import TextEmbedding
import os

# 设置API Key（从环境变量读取）
dashscope.api_key = os.getenv('DASHSCOPE_API_KEY')

# 生成文本向量
response = TextEmbedding.call(
    model='text-embedding-v2',  # 通义千问的Embedding模型
    input='人工智能正在改变世界'
)

# 获取向量
embedding = response.output['embeddings'][0]['embedding']
print(f"向量维度: {len(embedding)}")  # 1536维
print(f"向量前5个值: {embedding[:5]}")

from dashscope import TextEmbedding  # 从dashscope导入Embedding模型
import numpy as np                    # 导入numpy库，简写为np


def cosine_similarity(vec1, vec2):
    """计算两个向量的余弦相似度"""
    return np.dot(vec1, vec2) / (np.linalg.norm(vec1) * np.linalg.norm(vec2))


if __name__ == "__main__":
    # 生成两个文本的向量
    text1 = "机器学习是人工智能的一个分支"
    text2 = "深度学习属于机器学习领域"

    # 向量化文本1
    response1 = TextEmbedding.call(
        model='text-embedding-v2',
        input=text1
    )
    vec1 = response1.output['embeddings'][0]['embedding']

    # 向量化文本2
    response2 = TextEmbedding.call(
        model='text-embedding-v2',
        input=text2
    )
    vec2 = response2.output['embeddings'][0]['embedding']

    # 计算相似度
    similarity = cosine_similarity(vec1, vec2)
    print(f"'{text1}' 和 '{text2}' 的相似度: {similarity:.4f}")
    # 输出：相似度: 0.8523（非常相似！）

from dashscope import TextEmbedding
import numpy as np


def cosine_similarity(vec1, vec2):
    """计算两个向量的余弦相似度"""
    return np.dot(vec1, vec2) / (np.linalg.norm(vec1) * np.linalg.norm(vec2))


if __name__ == "__main__":
    # 批量向量化多个文本
    texts = [
        "人工智能",
        "机器学习",
        "深度学习",
        "神经网络"
    ]

    response = TextEmbedding.call(
        model='text-embedding-v2',
        input=texts  # 直接传入列表
    )

    # 获取所有向量
    embeddings = [item['embedding'] for item in response.output['embeddings']]
    print(f"成功向量化 {len(embeddings)} 个文本")

    # 计算所有文本之间的相似度矩阵
    for i, text1 in enumerate(texts):
        for j, text2 in enumerate(texts):
            if i < j:  # 只计算上三角
                sim = cosine_similarity(embeddings[i], embeddings[j])
                print(f"'{text1}' ↔ '{text2}': {sim:.4f}")

# 安装：pip install chromadb
import chromadb

# 1. 创建客户端
client = chromadb.Client()

# 2. 创建集合（使用默认embedding函数）
collection = client.create_collection("my_documents")

# 3. 添加文档（自动向量化，无需手动embedding）
collection.add(
    documents=[
        "Python是一种编程语言",
        "机器学习是人工智能的分支",
        "深度学习使用神经网络"
    ],
    ids=["doc1", "doc2", "doc3"]
)

# 4. 查询相似文档（自动embedding查询文本）
results = collection.query(
    query_texts=["什么是AI？"],
    n_results=2
)

print(results)

# 建议在虚拟环境中安装（venv/conda 都可以）

# 基础依赖
pip install chromadb dashscope

# 如果你后面需要读取 Excel（第1章加载Excel示例用到）
pip install pandas openpyxl

# macOS / Linux
export DASHSCOPE_API_KEY="你的key"

# Windows PowerShell
$env:DASHSCOPE_API_KEY="你的key"

import os
import dashscope
from dashscope import TextEmbedding


dashscope.api_key = os.getenv("DASHSCOPE_API_KEY")
if not dashscope.api_key:
    raise RuntimeError("请先设置环境变量 DASHSCOPE_API_KEY")


def embed_texts(texts):
    """输入：str 或 List[str]；输出：List[List[float]]"""
    if isinstance(texts, str):
        texts = [texts]

    resp = TextEmbedding.call(
        model="text-embedding-v2",
        input=texts
    )

    # 失败时给出明确报错，便于定位（Key/配额/网络）
    if not resp or not hasattr(resp, "output"):
        raise RuntimeError(f"Embedding 调用失败：{resp}")

    embeddings = [e["embedding"] for e in resp.output["embeddings"]]
    return embeddings


vec = embed_texts("人工智能正在改变世界")[0]
print(len(vec), vec[:8])

batch = embed_texts(["机器学习", "深度学习", "自然语言处理"])
print(len(batch), len(batch[0]))

import os
import chromadb
import dashscope
from dashscope import TextEmbedding


dashscope.api_key = os.getenv("DASHSCOPE_API_KEY")
if not dashscope.api_key:
    raise RuntimeError("请先设置环境变量 DASHSCOPE_API_KEY")


class DashScopeEmbeddingFunction:
    def __init__(self, model: str = "text-embedding-v2"):
        self.model = model

    def name(self) -> str:
        return f"dashscope:{self.model}"

    def _embed(self, texts):
        if isinstance(texts, str):
            texts = [texts]
        resp = TextEmbedding.call(model=self.model, input=texts)
        if not resp or not hasattr(resp, "output"):
            raise RuntimeError(f"Embedding 调用失败：{resp}")
        return [e["embedding"] for e in resp.output["embeddings"]]

    def embed_documents(self, input):
        return self._embed(input)

    def embed_query(self, input):
        return self._embed(input)

    def __call__(self, input):
        return self.embed_documents(input)


# 1) 创建客户端（开发阶段内存即可；生产建议 PersistentClient）
client = chromadb.Client()

# 2) 创建集合（关键：绑定 embedding_function）
collection = client.get_or_create_collection(
    name="my_documents",
    embedding_function=DashScopeEmbeddingFunction()
)

# 3) 写入（会自动调用 dashscope_embedding_function 做向量化）
collection.add(
    documents=[
        "人工智能（AI）是计算机科学的一个分支",
        "机器学习是人工智能的重要方向",
        "深度学习使用神经网络进行特征学习"
    ],
    ids=["doc1", "doc2", "doc3"],
    metadatas=[
        {"source": "demo", "type": "txt"},
        {"source": "demo", "type": "txt"},
        {"source": "demo", "type": "txt"}
    ]
)

# 4) 查询（同样会自动向量化 query_texts 并做相似检索）
results = collection.query(
    query_texts=["AI 包括哪些方向？"],
    n_results=2
)

print(results)

# 1. 持久化存储（数据保存到磁盘）
import chromadb

# 创建持久化客户端
client = chromadb.PersistentClient(path="./chroma_db")
collection = client.get_or_create_collection("my_docs")

# 2. 添加带元数据的文档
collection.add(
    documents=["Python是一种编程语言", "机器学习很有趣"],
    metadatas=[
        {"category": "编程", "difficulty": "初级"},
        {"category": "AI", "difficulty": "中级"}
    ],
    ids=["doc1", "doc2"]
)

# 3. 使用元数据过滤查询
results = collection.query(
    query_texts=["学习编程"],
    n_results=5,
    where={"category": "编程"}  # 只查询编程类文档
)

# 4. 更新文档
collection.update(
    ids=["doc1"],
    documents=["Python是一种强大的编程语言"]
)

# 5. 删除文档
collection.delete(ids=["doc2"])

# 完整RAG存储阶段实战示例
# 安装依赖：pip install "chromadb>=0.5" dashscope pypdf beautifulsoup4 pandas openpyxl

import chromadb
import dashscope
from dashscope import TextEmbedding
import pypdf
from bs4 import BeautifulSoup
import requests
import pandas as pd
from typing import List
import re
import os

# 配置API密钥
dashscope.api_key = os.getenv("DASHSCOPE_API_KEY", "")
if not dashscope.api_key:
    raise RuntimeError("请先设置环境变量 DASHSCOPE_API_KEY，然后再运行")


class RAGStoragePipeline:
    """RAG存储阶段完整流程"""

    def __init__(self, db_path: str = "./rag_db"):
        """初始化RAG存储管道"""
        self.db_path = db_path
        self.client = chromadb.PersistentClient(path=db_path)
        self.collection = None

    def load_documents(self, file_paths: List[str]) -> List[str]:
        """1️⃣ 加载文档"""
        print("🔄 开始加载文档...")
        documents = []

        for file_path in file_paths:
            if file_path.endswith('.pdf'):
                # 加载PDF文件
                content = self._load_pdf(file_path)
                documents.append(content)
                print(f"✅ 已加载PDF: {file_path}")

            elif file_path.endswith('.txt'):
                # 加载文本文件
                content = self._load_txt(file_path)
                documents.append(content)
                print(f"✅ 已加载TXT: {file_path}")

            elif file_path.startswith('http'):
                # 加载网页内容
                content = self._load_webpage(file_path)
                documents.append(content)
                print(f"✅ 已加载网页: {file_path}")

        print(f"🎉 共加载 {len(documents)} 个文档")
        return documents

    def _load_pdf(self, file_path: str) -> str:
        """加载PDF文件"""
        with open(file_path, 'rb') as file:
            reader = pypdf.PdfReader(file)
            text = ""
            for page in reader.pages:
                text += page.extract_text()
        return text

    def _load_txt(self, file_path: str) -> str:
        """加载文本文件"""
        with open(file_path, 'r', encoding='utf-8') as file:
            return file.read()

    def _load_webpage(self, url: str) -> str:
        """加载网页内容"""
        response = requests.get(url)
        soup = BeautifulSoup(response.content, 'html.parser')
        # 移除脚本和样式
        for script in soup(["script", "style"]):
            script.decompose()
        return soup.get_text()

    def split_documents(self, documents: List[str], chunk_size: int = 500, overlap: int = 50) -> List[str]:
        """2️⃣ 分割文档"""
        print("🔄 开始分割文档...")
        chunks = []

        for doc in documents:
            # 清理文本
            doc = re.sub(r'\s+', ' ', doc.strip())

            # 固定长度分割
            for i in range(0, len(doc), chunk_size - overlap):
                chunk = doc[i:i + chunk_size]
                if len(chunk.strip()) > 50:  # 过滤太短的片段
                    chunks.append(chunk.strip())

        print(f"🎉 共分割出 {len(chunks)} 个文档块")
        return chunks

    def create_embedding_function(self):
        """3️⃣ 创建向量化函数"""

        class DashScopeEmbeddingFunction:
            """符合 Chroma v0.5+ 的 EmbeddingFunction 协议（必须有 name()）"""

            def name(self):
                return "dashscope-text-embedding-v2"

            def _embed_texts(self, texts):
                if isinstance(texts, str):
                    texts = [texts]

                embeddings = []
                for text in texts:
                    response = TextEmbedding.call(
                        model='text-embedding-v2',
                        input=text
                    )
                    embeddings.append(response.output['embeddings'][0]['embedding'])
                return embeddings

            def __call__(self, input):
                return self._embed_texts(input)

            def embed_documents(self, documents):
                return self._embed_texts(documents)

            def embed_query(self, input):
                return self._embed_texts(input)

        return DashScopeEmbeddingFunction()

    def store_documents(self, chunks: List[str], collection_name: str = "rag_documents"):
        """4️⃣ 存储文档到向量数据库"""
        print("🔄 开始存储文档...")

        # 创建embedding函数
        embedding_function = self.create_embedding_function()

        # 创建集合
        # 注意：如果你之前用同名 collection 创建过（embedding_function 不同），Chroma 会做冲突校验并报错。
        # 解决办法：换一个 collection_name，或者删除/更换持久化目录（db_path）。
        self.collection = self.client.get_or_create_collection(
            name=collection_name,
            embedding_function=embedding_function
        )

        # 生成文档ID
        doc_ids = [f"doc_{i}" for i in range(len(chunks))]

        # 添加元数据
        metadatas = [
            {
                "source": f"chunk_{i}",
                "length": len(chunk),
                "created_at": "2024-01-01"
            }
            for i, chunk in enumerate(chunks)
        ]

        # 批量添加文档
        self.collection.add(
            documents=chunks,
            ids=doc_ids,
            metadatas=metadatas
        )

        print(f"🎉 成功存储 {len(chunks)} 个文档块到向量数据库")

    def run_complete_pipeline(self, file_paths: List[str], collection_name: str = "rag_documents"):
        """运行完整的RAG存储管道"""
        print("🚀 开始RAG存储阶段完整流程...")
        print("=" * 50)

        # 1️⃣ 加载文档
        documents = self.load_documents(file_paths)

        # 2️⃣ 分割文档
        chunks = self.split_documents(documents)

        # 3️⃣ 向量化 + 4️⃣ 存储
        self.store_documents(chunks, collection_name)

        print("=" * 50)
        print("🎉 RAG存储阶段完成！")
        print(f"📊 数据库路径: {self.db_path}")
        print(f"📋 集合名称: {collection_name}")
        print(f"📄 文档块数量: {len(chunks)}")

        return self.collection


# 使用示例
if __name__ == "__main__":
    # 创建RAG存储管道
    rag_pipeline = RAGStoragePipeline(db_path="./my_rag_db")

    # 准备文档路径（可以是PDF、TXT、网页URL）
    document_sources = [
        "../data/sample_document.pdf",  # 本地PDF文件
        "../data/article.txt",  # 本地文本文件
        "https://www.python.org",  # Python官网
    ]

    # 运行完整流程
    collection = rag_pipeline.run_complete_pipeline(
        file_paths=document_sources,
        collection_name="python_docs"
    )

    # 测试查询
    print("\n🔍 测试查询功能...")
    results = collection.query(
        query_texts=["什么是Python？"],
        n_results=3
    )

    print("查询结果:")
    for i, (doc, metadata, distance) in enumerate(zip(
            results['documents'][0],
            results['metadatas'][0],
            results['distances'][0]
    ), 1):
        print(f"\n{i}. 相似度: {1 - distance:.4f}")
        print(f"   文档: {doc[:100]}...")
        print(f"   元数据: {metadata}")