附录F：本地模拟生成式引擎环境（Docker Compose + Ollama）

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import ollama
import chromadb
from chromadb.utils import embedding_functions
import os
import logging

app = FastAPI(title="本地生成引擎模拟器")

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 环境变量
OLLAMA_BASE_URL = os.getenv("OLLAMA_BASE_URL", "http://ollama:11434")
CHROMA_HOST = os.getenv("CHROMA_HOST", "chromadb")
CHROMA_PORT = int(os.getenv("CHROMA_PORT", "8000"))
MODEL_NAME = os.getenv("MODEL_NAME", "qwen2.5:7b")
EMBEDDING_MODEL = os.getenv("EMBEDDING_MODEL", "bge-m3:latest")
COLLECTION_NAME = os.getenv("COLLECTION_NAME", "geo_test")

# 初始化 Ollama 客户端
ollama_client = ollama.Client(host=OLLAMA_BASE_URL)

# 初始化 ChromaDB 客户端
chroma_client = chromadb.HttpClient(host=CHROMA_HOST, port=CHROMA_PORT)
embedding_func = embedding_functions.OllamaEmbeddingFunction(
    model_name=EMBEDDING_MODEL,
    url=OLLAMA_BASE_URL
)

# 获取或创建集合
try:
    collection = chroma_client.get_collection(
        name=COLLECTION_NAME,
        embedding_function=embedding_func
    )
    logger.info(f"已加载现有集合: {COLLECTION_NAME}")
except:
    collection = chroma_client.create_collection(
        name=COLLECTION_NAME,
        embedding_function=embedding_func
    )
    logger.info(f"已创建新集合: {COLLECTION_NAME}")

class QueryRequest(BaseModel):
    query: str
    top_k: int = 5
    use_rag: bool = True
    system_prompt: str = "你是一个专业的搜索助手。请基于提供的上下文信息，用中文回答用户问题。如果信息不足，请明确说明。"

class IngestRequest(BaseModel):
    file_path: str = "/app/data"
    chunk_size: int = 500
    chunk_overlap: int = 50

@app.get("/health")
async def health():
    """健康检查"""
    return {"status": "ok", "model": MODEL_NAME, "collection": COLLECTION_NAME}

@app.post("/ingest")
async def ingest_content(request: IngestRequest):
    """导入内容到向量数据库"""
    import glob
    from langchain.text_splitter import RecursiveCharacterTextSplitter
    
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=request.chunk_size,
        chunk_overlap=request.chunk_overlap
    )
    
    files = glob.glob(f"{request.file_path}/**/*.md", recursive=True) + \
            glob.glob(f"{request.file_path}/**/*.txt", recursive=True) + \
            glob.glob(f"{request.file_path}/**/*.html", recursive=True)
    
    if not files:
        raise HTTPException(status_code=404, detail="未找到可导入的文件")
    
    imported_count = 0
    for file_path in files:
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                content = f.read()
            
            chunks = text_splitter.split_text(content)
            
            for i, chunk in enumerate(chunks):
                collection.add(
                    documents=[chunk],
                    metadatas=[{
                        "source": file_path,
                        "chunk_index": i,
                        "total_chunks": len(chunks)
                    }],
                    ids=[f"{file_path}_{i}"]
                )
                imported_count += 1
                
            logger.info(f"已导入: {file_path} ({len(chunks)} 块)")
        except Exception as e:
            logger.error(f"导入失败: {file_path} - {str(e)}")
    
    return {"imported_chunks": imported_count, "files_processed": len(files)}

@app.post("/query")
async def query(request: QueryRequest):
    """模拟生成引擎的问答流程"""
    # 步骤1: 检索相关上下文
    if request.use_rag:
        results = collection.query(
            query_texts=[request.query],
            n_results=request.top_k
        )
        
        if results['documents'] and results['documents'][0]:
            context = "\n\n".join([
                f"[来源: {meta['source']}]\n{doc}"
                for doc, meta in zip(results['documents'][0], results['metadatas'][0])
            ])
            
            sources = [
                {
                    "content": doc[:200] + "...",
                    "source": meta['source'],
                    "relevance": score
                }
                for doc, meta, score in zip(
                    results['documents'][0],
                    results['metadatas'][0],
                    results['distances'][0]
                )
            ]
        else:
            context = "未找到相关上下文。"
            sources = []
    else:
        context = ""
        sources = []
    
    # 步骤2: 构建 Prompt
    messages = [
        {"role": "system", "content": request.system_prompt},
        {"role": "user", "content": f"上下文信息：\n{context}\n\n用户问题：{request.query}"}
    ]
    
    # 步骤3: 调用 Ollama 生成回答
    try:
        response = ollama_client.chat(
            model=MODEL_NAME,
            messages=messages,
            options={
                "temperature": 0.3,
                "top_p": 0.9,
                "max_tokens": 2048
            }
        )
        
        return {
            "query": request.query,
            "answer": response['message']['content'],
            "sources": sources,
            "model": MODEL_NAME,
            "rag_enabled": request.use_rag
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"模型推理失败: {str(e)}")

@app.post("/compare")
async def compare_answers(request: QueryRequest):
    """对比有无 RAG 的回答差异"""
    # 有 RAG
    rag_response = await query(QueryRequest(
        query=request.query,
        use_rag=True,
        system_prompt=request.system_prompt
    ))
    
    # 无 RAG（纯模型知识）
    no_rag_response = await query(QueryRequest(
        query=request.query,
        use_rag=False,
        system_prompt=request.system_prompt
    ))
    
    return {
        "query": request.query,
        "with_rag": rag_response,
        "without_rag": no_rag_response
    }

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8080)

# batch_test.py
import requests
import json
import time

BASE_URL = "http://localhost:8080"

test_queries = [
    "你的产品的主要功能是什么？",
    "你的产品与竞争对手相比有什么优势？",
    "你的产品如何安装？",
    "你的产品的价格是多少？",
    "你的产品支持哪些平台？"
]

results = []
for query in test_queries:
    response = requests.post(
        f"{BASE_URL}/query",
        json={"query": query, "use_rag": True}
    )
    result = response.json()
    results.append({
        "query": query,
        "answer": result["answer"][:200],
        "sources_count": len(result["sources"])
    })
    print(f"Query: {query}")
    print(f"Answer: {result['answer'][:100]}...")
    print(f"Sources: {len(result['sources'])}")
    print("-" * 50)
    time.sleep(1)

# 保存结果
with open("geo_test_results.json", "w", encoding="utf-8") as f:
    json.dump(results, f, ensure_ascii=False, indent=2)