#RAG 后端系统集成：FastAPI 与向量数据库的联动实践

import chromadb
from chromadb.config import Settings
from sentence_transformers import SentenceTransformer

class ChromaStore:
    def __init__(self, collection_name: str = "docs"):
        # 关闭匿名遥测，本地持久化存储
        self.client = chromadb.PersistentClient(
            path="./chroma_db",
            settings=Settings(anonymized_telemetry=False)
        )
        # 免费轻量的英文通用嵌入模型（中文可选shibing624/text2vec-base-chinese）
        self.embedding_model = SentenceTransformer("all-MiniLM-L6-v2")
        # 用余弦相似度创建/获取集合
        self.collection = self.client.get_or_create_collection(
            name=collection_name,
            metadata={"hnsw:space": "cosine"}
        )

    def add_doc(self, doc_id: str, text: str, metadata: dict | None = None):
        """单条文档入库"""
        embedding = self.embedding_model.encode(text).tolist()
        self.collection.add(
            ids=[doc_id],
            documents=[text],
            embeddings=[embedding],
            metadatas=[metadata or {}]
        )

    def search(self, query: str, top_k: int = 5) -> list[dict]:
        """语义检索，返回格式统一的结果"""
        query_emb = self.embedding_model.encode(query).tolist()
        raw_res = self.collection.query(
            query_embeddings=[query_emb],
            n_results=top_k,
            include=["documents", "metadatas", "distances"]
        )
        # 把距离（越小越相似）转成直观的相似度（0-1）
        return [
            {
                "id": raw_res["ids"][0][i],
                "text": raw_res["documents"][0][i],
                "metadata": raw_res["metadatas"][0][i],
                "score": round(1 - raw_res["distances"][0][i], 4)
            }
            for i in range(len(raw_res["ids"][0]))
        ]

    def delete_doc(self, doc_id: str):
        self.collection.delete(ids=[doc_id])

# 全局实例，服务启动时加载
vector_db = ChromaStore()

from pymilvus import MilvusClient, DataType

class MilvusStore:
    def __init__(self, uri: str = "http://localhost:19530"):
        self.client = MilvusClient(uri=uri)
        self.collection_name = "docs"
        self.embedding_dim = 384  # 对应all-MiniLM-L6-v2的输出维度
        # 服务启动时自动创建不存在的集合
        if not self.client.has_collection(self.collection_name):
            self._create_collection()

    def _create_collection(self):
        schema = MilvusClient.create_schema(
            auto_id=True,
            enable_dynamic_field=True  # 允许存任意额外metadata
        )
        schema.add_field("id", DataType.INT64, is_primary=True)
        schema.add_field("vector", DataType.FLOAT_VECTOR, dim=self.embedding_dim)
        schema.add_field("text", DataType.VARCHAR, max_length=4096)
        schema.add_field("source", DataType.VARCHAR, max_length=255)

        self.client.create_collection(
            collection_name=self.collection_name,
            schema=schema,
            index_params={
                "vector": {"type": "HNSW", "metric_type": "COSINE", "params": {"M": 8, "efConstruction": 64}}
            }
        )
        # 加载到内存提高检索速度
        self.client.load_collection(self.collection_name)

    def batch_insert(self, texts: list[str], sources: list[str], embeddings: list[list[float]]):
        """批量文档入库（单条慢，生产必须用批量）"""
        data = [
            {"vector": emb, "text": text, "source": source}
            for text, source, emb in zip(texts, sources, embeddings)
        ]
        return self.client.insert(collection_name=self.collection_name, data=data)

# processing/chunker.py
from typing import List
import re

class TextChunker:
    def __init__(self, chunk_size: int = 500, overlap: int = 50):
        self.chunk_size = chunk_size  # 每块最多字符数（英文≈80token，中文≈250token）
        self.overlap = overlap        # 相邻块的重叠字符数，保留上下文

    def chunk_text(self, text: str, source: str = "") -> List[dict]:
        """优先按段落分割，不够再补，超了再切"""
        # 先清理多余空行
        paragraphs = [p.strip() for p in re.split(r'\n{2,}', text) if p.strip()]
        chunks = []
        current_chunk = ""
        chunk_idx = 0

        for para in paragraphs:
            # 当前块+下一段没超，直接加
            if len(current_chunk) + len(para) + 2 < self.chunk_size:
                current_chunk += para + "\n\n"
            else:
                # 当前块不为空，先存
                if current_chunk.strip():
                    chunks.append({
                        "id": f"{source.replace('/', '_')}_{chunk_idx}",
                        "text": current_chunk.strip(),
                        "metadata": {"source": source, "chunk_idx": chunk_idx}
                    })
                    chunk_idx += 1
                # 滑动窗口初始化：下一段的最后overlap字符作为开头（可选）
                if len(para) > self.overlap:
                    current_chunk = para[-self.overlap:] + "\n\n" + para + "\n\n"
                else:
                    current_chunk = para + "\n\n"

        # 最后补存剩余的块
        if current_chunk.strip():
            chunks.append({
                "id": f"{source.replace('/', '_')}_{chunk_idx}",
                "text": current_chunk.strip(),
                "metadata": {"source": source, "chunk_idx": chunk_idx}
            })
        return chunks

from vector_store.chroma_store import vector_db
from sentence_transformers import SentenceTransformer
from openai import AsyncOpenAI
import os
from dotenv import load_dotenv

# 加载环境变量（API Key不要写在代码里！）
load_dotenv()

class RAGService:
    def __init__(self):
        self.embedding_model = SentenceTransformer("all-MiniLM-L6-v2")
        # 支持OpenAI兼容接口（如Ollama、DeepSeek、通义千问）
        self.llm_client = AsyncOpenAI(
            api_key=os.getenv("OPENAI_API_KEY"),
            base_url=os.getenv("OPENAI_BASE_URL", "https://api.openai.com/v1")
        )

    async def _retrieve(self, query: str, top_k: int = 5, min_score: float = 0.4) -> list[dict]:
        """内部检索方法，过滤掉低相似度的噪声"""
        results = vector_db.search(query, top_k=top_k)
        return [r for r in results if r["score"] >= min_score]

    async def generate_answer(self, query: str, top_k: int = 5, min_score: float = 0.4, stream: bool = False):
        """对外暴露的检索+生成方法"""
        # 1. 检索
        docs = await self._retrieve(query, top_k, min_score)
        if not docs:
            return {"answer": "抱歉，知识库中未找到相关内容，请尝试换个说法。", "sources": []}

        # 2. 构建增强Prompt
        context = "\n\n".join([
            f"[来源 {i+1}]({doc['metadata']['source']}, 块号 {doc['metadata']['chunk_idx']})\n{doc['text']}"
            for i, doc in enumerate(docs)
        ])
        prompt = f"""你是一个严谨的知识库助手，请严格遵循以下规则回答：
1. 仅参考【参考资料】回答，资料外的内容不要编造
2. 答案要简洁明了
3. 如果资料中没有明确答案，请直接说「未找到相关信息」

参考资料：
{context}

用户问题：{query}"""

        # 3. 调用LLM生成
        llm_response = await self.llm_client.chat.completions.create(
            model=os.getenv("LLM_MODEL", "gpt-4o-mini"),
            messages=[{"role": "user", "content": prompt}],
            temperature=0.2,  # 低温度保证稳定
            max_tokens=1000,
            stream=stream
        )
        return {"response": llm_response, "sources": docs}

from fastapi import APIRouter, HTTPException
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
from services.rag_service import RAGService
from processing.chunker import TextChunker
from vector_store.chroma_store import vector_db
from sentence_transformers import SentenceTransformer
import json
import asyncio

router = APIRouter(prefix="/api/rag", tags=["RAG 接口"])
rag_service = RAGService()
chunker = TextChunker()
embedding_model = SentenceTransformer("all-MiniLM-L6-v2")

# 请求/响应模型定义
class QueryReq(BaseModel):
    query: str
    top_k: int = 5
    min_score: float = 0.4
    stream: bool = False

class IngestReq(BaseModel):
    text: str
    source: str = "manual_input"

# ──────────────────────────────────────
# 普通问答接口
# ──────────────────────────────────────
@router.post("/query")
async def rag_query(req: QueryReq):
    result = await rag_service.generate_answer(**req.model_dump())
    if not req.stream:
        return {
            "answer": result["response"].choices[0].message.content,
            "sources": [{"text": d["text"][:200] + "...", "source": d["metadata"]["source"], "score": d["score"]} for d in result["sources"]]
        }

# ──────────────────────────────────────
# 流式问答接口（打字机效果）
# ──────────────────────────────────────
@router.post("/query/stream")
async def rag_query_stream(req: QueryReq):
    result = await rag_service.generate_answer(**req.model_dump())
    if not result["sources"]:
        return {"answer": "未找到相关内容", "sources": []}

    async def event_generator():
        # 先推送来源（可选）
        yield f"data: {json.dumps({'type': 'sources', 'content': result['sources']})}\n\n"
        # 推送LLM token流
        async for chunk in result["response"]:
            if token := chunk.choices[0].delta.content:
                yield f"data: {json.dumps({'type': 'token', 'content': token})}\n\n"
        # 推送结束标记
        yield f"data: {json.dumps({'type': 'done'})}\n\n"

    return StreamingResponse(
        event_generator(),
        media_type="text/event-stream",
        headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"}
    )

# ──────────────────────────────────────
# 文档手动入库接口（测试用）
# ──────────────────────────────────────
@router.post("/ingest")
async def ingest(req: IngestReq):
    chunks = chunker.chunk_text(req.text, req.source)
    for chunk in chunks:
        vector_db.add_doc(**chunk)
    return {"status": "success", "chunks_count": len(chunks)}