from pymilvus import model

# 使用 OpenAI 嵌入模型
embed_fn = model.dense.OpenAIEmbeddingFunction(
    model_name="text-embedding-3-small",
    api_key="your-api-key"
)

# 将文本转换为向量
texts = [
    "Milvus 是一个开源的向量数据库",
    "Claude 是 Anthropic 开发的 AI 助手",
    "向量搜索在 AI 应用中非常重要"
]

vectors = embed_fn.encode_documents(texts)

# 每个向量是 1536 维的浮点数组
print(f"向量维度: {len(vectors[0])}")  # 1536
print(f"向量示例: {vectors[0][:5]}")   # [0.012, -0.034, ...]

from pymilvus import MilvusClient, DataType

client = MilvusClient("./milvus_demo.db")

# 创建集合 Schema
schema = client.create_schema(auto_id=True)
schema.add_field("id", DataType.INT64, is_primary=True)
schema.add_field("vector", DataType.FLOAT_VECTOR, dim=1536)
schema.add_field("text", DataType.VARCHAR, max_length=65535)
schema.add_field("source", DataType.VARCHAR, max_length=256)

# 创建集合
client.create_collection(
    collection_name="knowledge_base",
    schema=schema,
)

# 插入数据
data = [{
    "vector": vectors[0],
    "text": "Milvus 是一个开源的向量数据库",
    "source": "milvus_docs.md"
}]
client.insert("knowledge_base", data)

from pymilvus import MilvusClient, model
from langchain.text_splitter import RecursiveCharacterTextSplitter

class DocumentIngestion:
    """文档摄入管道"""

    def __init__(self, milvus_client: MilvusClient):
        self.client = milvus_client
        self.embed_fn = model.dense.OpenAIEmbeddingFunction(
            model_name="text-embedding-3-small"
        )
        self.splitter = RecursiveCharacterTextSplitter(
            chunk_size=512,
            chunk_overlap=50,
            separators=["\n\n", "\n", "。", ".", " "]
        )

    def ingest(self, documents: list[dict]) -> int:
        """摄入一批文档，返回插入的块数"""
        chunks = []
        for doc in documents:
            # 1. 分块
            text_chunks = self.splitter.split_text(doc["content"])

            # 2. 嵌入
            vectors = self.embed_fn.encode_documents(text_chunks)

            # 3. 组装数据
            for i, (text, vec) in enumerate(
                zip(text_chunks, vectors)
            ):
                chunks.append({
                    "vector": vec,
                    "text": text,
                    "source": doc["source"],
                    "chunk_index": i
                })

        # 4. 批量插入
        self.client.insert("knowledge_base", chunks)
        return len(chunks)

import anthropic

class RAGQueryEngine:
    """RAG 查询引擎"""

    def __init__(self, milvus_client, embed_fn):
        self.milvus = milvus_client
        self.embed_fn = embed_fn
        self.claude = anthropic.Anthropic()

    def query(self, question: str, top_k: int = 5) -> dict:
        # 1. 将问题转化为向量
        query_vec = self.embed_fn.encode_queries([question])[0]

        # 2. 在 Milvus 中检索最相关的文档块
        results = self.milvus.search(
            collection_name="knowledge_base",
            data=[query_vec],
            limit=top_k,
            output_fields=["text", "source"]
        )

        # 3. 组装上下文
        context = "\n\n".join([
            hit["entity"]["text"]]
            for hit in results[0]
        ])

        # 4. 调用 Claude 生成回答
        response = self.claude.messages.create(
            model="claude-sonnet-4-20250514",
            max_tokens=2048,
            system="基于以下上下文回答问题。如果上下文中没有相关信息，请说明。",
            messages=[{"role": "user", "content": f"上下文：{context}\n\n问题：{question}"}]
        )

        return {
            "answer": response.content[0].text,
            "sources": [hit["entity"]["source"] for hit in results[0]]
        }

from pymilvus import AnnSearchRequest, WeightedRanker

# 向量搜索请求
vector_search = AnnSearchRequest(
    data=[query_vector],
    anns_field="vector",
    param={"metric_type": "COSINE", "params": {"nprobe": 10}},
    limit=20
)

# 全文搜索请求（Milvus 2.4+）
text_search = AnnSearchRequest(
    data=[query_sparse_vector],
    anns_field="sparse_vector",
    param={"metric_type": "IP"},
    limit=20
)

# 混合搜索：加权融合两种结果
results = client.hybrid_search(
    collection_name="knowledge_base",
    reqs=[vector_search, text_search],
    ranker=WeightedRanker(0.7, 0.3),  # 向量70% + 关键词30%
    limit=10
)

class KnowledgeQA:
    """知识库问答系统 — 请补充完整实现"""

    def __init__(self):
        # TODO: 初始化 Milvus 客户端和嵌入模型
        pass

    # TODO 1: 实现文档摄入
    def ingest_docs(self, folder_path: str) -> int:
        # 提示：遍历文件夹，分块，嵌入，插入
        pass

    # TODO 2: 实现混合检索
    def search(self, query: str, top_k: int = 5) -> list:
        # 提示：查询嵌入 + Milvus search
        pass

    # TODO 3: 实现 RAG 生成
    def answer(self, question: str) -> dict:
        # 提示：search → 组装上下文 → Claude 生成
pass