How I Used Gemini CLI to Orchestrate a Complex RAG Migration

# 旧版架构（LangChain 0.0.x）
from langchain import OpenAI, VectorDBQA
from pgvector.sqlalchemy import Vector
from sqlalchemy import create_engine

class OldRAGSystem:
    def __init__(self):
        self.llm = OpenAI(temperature=0)
        self.engine = create_engine("postgresql://localhost:5432/vectors")
        self.vectorstore = VectorDBQA.from_connector(
            name="docs",
            description="Company documentation",
            embedding_function=self._get_embedding,
            engine=self.engine,
            table_name="document_vectors"
        )

    def query(self, question: str) -> str:
        # 简单直接：检索 -> 直接返回
        docs = self.vectorstore.similarity_search(question, k=5)
        context = "\n".join([d.page_content for d in docs])
        return self.llm.call(f"Context: {context}\n\nQuestion: {question}")

Phase 1: 现状分析（Gemini 完成）
Phase 2: 基础设施搭建（Gemini + 人工）
Phase 3: 核心逻辑迁移（Gemini 完成）
Phase 4: 数据迁移（人工 + Gemini 辅助脚本）
Phase 5: 灰度验证（人工 + Gemini 辅助）

# 启动 Gemini CLI 分析模式
gemini-cli chat --model gemini-2.5-pro

# 我的输入：
我需要你分析整个 /path/to/old_rag 目录的代码结构。
请输出：
1. 代码模块划分（每个模块的功能）
2. API 调用链路（从 query 到最终返回的完整流程）
3. 外部依赖（有哪些外部服务、数据库、API）
4. 配置管理（hardcode 的配置有哪些）
5. 测试覆盖（有哪些测试，测试覆盖了哪些场景）

## 代码模块划分

### 模块 1: vector_storage（向量存储层）
- `embedding_service.py`: 文本向量化
- `pgvector_store.py`: pgvector 存储和检索
- `chunking_strategy.py`: 文档分块策略

问题：embedding_service 使用的是 OpenAI ada-002，hardcode

### 模块 2: retrieval（检索层）
- `semantic_search.py`: 语义检索
- `keyword_search.py`: 关键词检索
- `hybrid_search.py`: 混合检索（但实现有问题）

问题：混合检索是简单 OR，没有真正融合 score

### 模块 3: generation（生成层）
- `llm_service.py`: LLM 调用封装
- `prompt_template.py`: Prompt 管理

问题：Prompt 散落在代码各处，没有统一管理

### 模块 4: api（对外接口）
- `fastapi_server.py`: API 服务
- `endpoints/`: 各端点定义

### 外部依赖
- PostgreSQL + pgvector
- OpenAI API（ ada-002 + GPT-4）
- Redis（缓存）

### Hardcode 配置
- embedding model: text-embedding-ada-002
- chunk_size: 512
- top_k: 5
- temperature: 0

# 新的 embedding 配置
EMBEDDING_CONFIG = {
    "provider": "openai",
    "model": "text-embedding-3-small",  # 从 ada-002 升级
    "dimension": 1536,  # 从 1536 降到 1536（3-small 支持 256-3072）
    "batch_size": 100,
}

# Gemini 生成的 Pinecone 存储层代码
from pinecone import Pinecone, ServerlessSpec
from langchain_pinecone import PineconeVectorStore
from langchain_openai import OpenAIEmbeddings

class NewVectorStore:
    def __init__(self, api_key: str, index_name: str):
        self.pc = Pinecone(api_key=api_key)

        # 确保索引存在
        if index_name not in [idx.name for idx in self.pc.list_indexes()]:
            self.pc.create_index(
                name=index_name,
                dimension=1536,
                metric="cosine",
                spec=ServerlessSpec(cloud="aws", region="us-east-1")
            )

        self.index = self.pc.Index(index_name)
        self.embeddings = OpenAIEmbeddings(
            model="text-embedding-3-small"
        )
        self.vectorstore = PineconeVectorStore(
            index=self.index,
            embedding=self.embeddings,
            text_key="text"
        )

    def add_documents(self, documents: list, ids: list = None):
        """批量添加文档"""
        return self.vectorstore.add_documents(documents=documents, ids=ids)

    def similarity_search(self, query: str, k: int = 5, filter: dict = None):
        """相似度检索"""
        return self.vectorstore.similarity_search(
            query=query,
            k=k,
            filter=filter
        )

    def similarity_search_with_score(self, query: str, k: int = 5):
        """带分数的检索"""
        return self.vectorstore.similarity_search_with_score(
            query=query,
            k=k
        )

# Gemini 生成的混合检索实现
from langchain.retrievers import EnsembleRetriever
from langchain_community.retrievers import BM25Retriever
from langchain_core.documents import Document
from rank_bm25 import BM25Okapi
import numpy as np

class HybridSearchRetriever:
    def __init__(self, vectorstore, documents: list):
        self.vectorstore = vectorstore

        # 构建 BM25 检索器
        self.doc_texts = [doc.page_content for doc in documents]
        self.bm25 = BM25Okapi(self.doc_texts)

    def retrieve(self, query: str, k: int = 10):
        """
        混合检索：向量检索 + BM25，RRF 融合
        RRF (Reciprocal Rank Fusion): 1/(rank + k)，k=60 是经验值
        """
        # 向量检索
        vector_results = self.vectorstore.similarity_search_with_score(query, k=k)
        vector_docs = [doc for doc, score in vector_results]

        # BM25 检索
        tokenized_query = query.split()
        bm25_scores = self.bm25.get_scores(tokenized_query)
        top_k_indices = np.argsort(bm25_scores)[::-1][:k]
        bm25_docs = [Document(page_content=self.doc_texts[i])
                     for i in top_k_indices]

        # RRF 融合
        k_rrf = 60  # RRF 参数
        rrf_scores = {}

        for rank, (doc, score) in enumerate(vector_results):
            doc_id = doc.page_content[:100]  # 用前 100 字符做 ID
            rrf_scores[doc_id] = rrf_scores.get(doc_id, 0) + 1 / (rank + k_rrf)

        for rank, doc in enumerate(bm25_docs):
            doc_id = doc.page_content[:100]
            rrf_scores[doc_id] = rrf_scores.get(doc_id, 0) + 1 / (rank + k_rrf)

        # 排序返回
        sorted_docs = sorted(rrf_scores.items(), key=lambda x: x[1], reverse=True)
        final_docs = [Document(page_content=doc_id) for doc_id, _ in sorted_docs[:k]]

        return final_docs

# 修正后的实现
def retrieve(self, query: str, k: int = 10, doc_ids: list = None):
    """
    混合检索 RRF 融合
    doc_ids: 文档唯一 ID 列表
    """
    # ... 向量检索和 BM25 检索同上 ...

    # 用 doc_id 而不是 page_content 前 100 字符
    for rank, (doc, score) in enumerate(vector_results):
        doc_id = doc.metadata.get('id') or doc_ids[rank]
        rrf_scores[doc_id] = rrf_scores.get(doc_id, 0) + 1 / (rank + k_rrf)

    for rank, (doc, score) in enumerate(bm25_scores[:k]):
        doc_id = doc_ids[rank]
        rrf_scores[doc_id] = rrf_scores.get(doc_id, 0) + 1 / (rank + k_rrf)

# Gemini 生成的 LLM 服务封装
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough

class NewLLMService:
    def __init__(self, model: str = "gpt-4o", temperature: float = 0):
        self.llm = ChatOpenAI(
            model=model,
            temperature=temperature
        )

    def create_chain(self, system_prompt: str, retrieval_chain):
        """创建 RAG chain"""
        prompt = ChatPromptTemplate.from_messages([
            ("system", system_prompt),
            ("human", "{question}")
        ])

        chain = (
            {"context": retrieval_chain, "question": RunnablePassthrough()}
            | prompt
            | self.llm
            | StrOutputParser()
        )

        return chain

    def query(self, question: str, context: str) -> str:
        """直接调用（不走 retrieval chain）"""
        prompt = f"""Given the following context:
{context}

Answer the question:
{question}
"""
        return self.llm.invoke(prompt)

# 数据迁移脚本
from tqdm import tqdm

class VectorMigrator:
    def __init__(self, old_pg_conn, new_pinecone_index):
        self.old_pg = old_pg_conn
        self.new_index = new_pinecone_index

    def migrate_batch(self, batch_size: int = 1000):
        """
        批量迁移向量数据
        """
        # 从 pgvector 读取
        query = f"""
        SELECT id, content, embedding, metadata
        FROM document_vectors
        WHERE migrated = FALSE
        ORDER BY id
        LIMIT {batch_size}
        """

        rows = self.old_pg.execute(query).fetchall()

        if not rows:
            print("No more rows to migrate")
            return 0

        # 转换成 Pinecone 格式
        vectors = []
        ids_to_update = []

        for row in rows:
            vectors.append({
                "id": str(row.id),
                "values": row.embedding,  # pgvector 是 list 类型
                "metadata": {
                    "content": row.content[:2000],  # Pinecone metadata 有 40KB 限制
                    **row.metadata
                }
            })
            ids_to_update.append(row.id)

        # 写入 Pinecone
        self.new_index.upsert(vectors)

        # 标记已迁移
        update_query = """
        UPDATE document_vectors
        SET migrated = TRUE, migrated_at = NOW()
        WHERE id = ANY(:ids)
        """
        self.old_pg.execute(update_query, {"ids": ids_to_update})

        return len(rows)

    def migrate_all(self, total: int):
        """
        完整迁移
        """
        migrated = 0
        batch_size = 1000

        with tqdm(total=total) as pbar:
            while True:
                count = self.migrate_batch(batch_size)
                if count == 0:
                    break
                migrated += count
                pbar.update(count)

        print(f"Migration complete: {migrated} vectors")

# 使用
migrator = VectorMigrator(old_pg, pinecone_index)
migrator.migrate_all(total=500000)

# 验证脚本
import json
from datetime import datetime

class RAGValidator:
    def __init__(self, old_rag, new_rag, test_queries: list):
        self.old_rag = old_rag
        self.new_rag = new_rag
        self.test_queries = test_queries
        self.results = []

    def run_validation(self):
        """运行验证测试"""
        for query in self.test_queries:
            old_result = self.old_rag.query(query)
            new_result = self.new_rag.query(query)

            # 用 embedding similarity 比较两个回答的语义相似度
            old_embedding = get_embedding(old_result)
            new_embedding = get_embedding(new_result)
            similarity = cosine_similarity(old_embedding, new_embedding)

            self.results.append({
                "query": query,
                "old_result": old_result,
                "new_result": new_result,
                "similarity": similarity,
                "passed": similarity > 0.85  # 阈值 0.85
            })

        return self.results

    def generate_report(self) -> dict:
        """生成验证报告"""
        passed = sum(1 for r in self.results if r["passed"])
        total = len(self.results)

        return {
            "timestamp": datetime.now().isoformat(),
            "total_queries": total,
            "passed": passed,
            "failed": total - passed,
            "pass_rate": f"{passed/total*100:.1f}%",
            "failed_queries": [r for r in self.results if not r["passed"]]
        }