QMDSearch/scripts/smoke_live.py

"""真实 Ollama 现场冒烟脚本

用法：uv run python scripts/smoke_live.py

链路：内存 Qdrant + 真实本地 Ollama（deepseek-r1:8b）+ 确定性哈希向量（无语义，仅打通流程）。
读取回归集第一篇文档走完整 ingest → search 流程，打印各阶段结果、原始 LLM 输出与耗时。

deepseek-r1 是推理模型，/api/generate 的 response 字段可能带 <think> 思考内容；
classifier / query_parser 内建 JSON 容错，若解析失败会走兜底路径——均为合理的现场观察。
"""

import asyncio
import json
import random
import sys
import time
from hashlib import sha256
from pathlib import Path

# 脚本直接运行时需要把项目根目录加入 sys.path
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from qdrant_client import AsyncQdrantClient  # noqa: E402

from app.config import settings  # noqa: E402
from app.core.chunker import Chunker  # noqa: E402
from app.core.classifier import Classifier  # noqa: E402
from app.core.ingestion import Ingester  # noqa: E402
from app.core.query_parser import QueryParser  # noqa: E402
from app.core.retriever import Retriever  # noqa: E402
from app.core.sparse import SparseEncoder  # noqa: E402
from app.core.summarizer import Summarizer  # noqa: E402
from app.models.document import DocumentInput  # noqa: E402
from app.models.knowledge import load_taxonomy  # noqa: E402
from app.models.search import SearchRequest  # noqa: E402
from app.services.ollama import OllamaClient  # noqa: E402
from app.services.qdrant import QdrantService  # noqa: E402

OLLAMA_MODEL = "deepseek-r1:8b"
OLLAMA_TIMEOUT = 600.0


class DeterministicEmbedding:
    """稳定哈希伪向量：同文本恒同向量，维度等于 settings.embedding_dimension（无语义）"""

    async def embed(self, texts: list[str]) -> list[list[float]]:
        vectors: list[list[float]] = []
        for text in texts:
            seed = int.from_bytes(sha256(text.encode("utf-8")).digest()[:8], "big")
            rng = random.Random(seed)
            vectors.append([rng.random() for _ in range(settings.embedding_dimension)])
        return vectors


class RecordingOllama:
    """包装真实 OllamaClient，记录每次调用的原始输出供现场观察"""

    def __init__(self, inner: OllamaClient) -> None:
        self.inner = inner
        self.records: list[tuple[str, str]] = []  # (调用用途, 原始输出)

    async def generate(self, prompt: str, json_mode: bool = False) -> str:
        kind = _prompt_kind(prompt)
        raw = await self.inner.generate(prompt, json_mode=json_mode)
        self.records.append((kind, raw))
        return raw


def _prompt_kind(prompt: str) -> str:
    """按 prompt 特征串标注调用用途"""
    if "你是知识库分类助手" in prompt:
        return "文档分类"
    if "你是搜索查询分析助手" in prompt:
        return "query 解析"
    if "请用一句话对以下文档内容进行高度概括" in prompt:
        return "L1 总结"
    if "请提取以下文档的主要章节结构" in prompt:
        return "L2 大纲"
    if "请对以下文档的每个章节/主题进行详细的内容摘要" in prompt:
        return "L3 内容大纲"
    if "请对以下文档内容进行详细摘要" in prompt:
        return "L2.5 摘要"
    return "未知"


def _fmt(seconds: float) -> str:
    return f"{seconds:.1f}s"


async def main() -> int:
    started = time.perf_counter()

    # 0. 环境检查
    ollama_inner = OllamaClient(base_url="http://localhost:11434", model=OLLAMA_MODEL, timeout=OLLAMA_TIMEOUT)
    if not await ollama_inner.is_available():
        print("Ollama 不可用（http://localhost:11434），请先启动 Ollama 服务")
        return 1
    print(f"[环境] Ollama 可用，模型 {OLLAMA_MODEL}；Qdrant 使用 :memory: 本地模式")

    # 1. 读取回归集第一篇文档与一条相关 query
    regression_path = Path(__file__).resolve().parent / "eval" / "regression_set.json"
    regression = json.loads(regression_path.read_text(encoding="utf-8"))
    doc_data = regression["documents"][0]
    query = next(q["query"] for q in regression["queries"] if q.get("golden_doc_id") == doc_data["id"])
    doc = DocumentInput(text=doc_data["text"], title=doc_data["title"])
    print(f"[数据] 文档《{doc.title}》（{len(doc.text)} 字符）；query：{query}")

    # 2. 组装真实链路（仅 embedding 为确定性伪向量）
    qdrant = QdrantService(client=AsyncQdrantClient(location=":memory:"))
    await qdrant.ensure_collections()
    ollama = RecordingOllama(ollama_inner)
    taxonomy = load_taxonomy()
    embedding = DeterministicEmbedding()
    ingester = Ingester(
        summarizer=Summarizer(ollama=ollama),  # type: ignore[arg-type]
        classifier=Classifier(ollama=ollama, taxonomy=taxonomy),  # type: ignore[arg-type]
        chunker=Chunker(),
        embedding=embedding,
        sparse=SparseEncoder(),
        qdrant=qdrant,
    )
    retriever = Retriever(
        qdrant=qdrant,
        query_parser=QueryParser(ollama=ollama, taxonomy=taxonomy),  # type: ignore[arg-type]
        embedding=embedding,
        sparse_encoder=SparseEncoder(),
    )

    # 3. 入库
    t0 = time.perf_counter()
    result = await ingester.ingest(doc)
    t_ingest = time.perf_counter() - t0
    print(f"\n[入库] 耗时 {_fmt(t_ingest)}")
    print(f"  总结层级: {result.summary.level.value}")
    print(f"  L1 总结: {result.summary.l1_summary[:100]}")
    print(f"  L2 大纲: {(result.summary.l2_outline or '(None，2.5 级回退)')[:100]}")
    print(f"  分类: {result.category} (置信度 {result.category_confidence:.2f})，tags={result.tags}")
    print(f"  chunks 数: {result.chunks_count}")

    # 4. 检索
    t0 = time.perf_counter()
    response = await retriever.search(SearchRequest(query=query))
    t_search = time.perf_counter() - t0
    print(f"\n[检索] 耗时 {_fmt(t_search)}")
    print(f"  routed_categories={response.routed_categories}，fallback={response.fallback}")
    print(f"  hits 数: {len(response.hits)}")
    if response.hits:
        hit = response.hits[0]
        print(f"  首条 hit: section_path={hit.section_path!r}，score={hit.score:.4f}")
        print(f"  首条 hit.doc_summary 前 50 字: {hit.doc_summary[:50]!r}")
        print(f"  首条 hit.text 前 50 字: {hit.text[:50]!r}")

    # 5. r1 模型原始输出观察（thinking / JSON 表现）
    print("\n[Ollama 原始输出摘录]（r1 推理模型的 JSON 表现现场观察）")
    for kind, raw in ollama.records:
        excerpt = raw.strip().replace("\n", " ")[:200]
        has_think = "<think>" in raw
        print(f"  - {kind}: 长度 {len(raw)}，含 <think>={has_think}，输出摘录: {excerpt!r}")

    print(f"\n[总耗时] {_fmt(time.perf_counter() - started)}")
    return 0


if __name__ == "__main__":
    sys.exit(asyncio.run(main()))