初始提交：餐智库 CIBank 餐饮行业知识库

backend/rag.py

@@ -0,0 +1,333 @@
+#!/usr/bin/env python3
+"""RAG 问答模块：向量检索 + LLM 生成。
+
+流程：
+1. 将用户问题转为向量
+2. 在 embeddings 表中做余弦相似度检索，取 TOP-K 相关文章
+3. 将检索到的文章内容作为上下文，调用千问生成回答
+4. 返回回答 + 引用来源（可追溯）
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import sqlite3
+import urllib.error
+import urllib.request
+from pathlib import Path
+
+import numpy as np
+
+from dotenv import load_dotenv
+from structurer import get_api_key, get_embeddings, DASHSCOPE_CHAT_URL
+
+BASE_DIR = Path(__file__).resolve().parent
+load_dotenv(BASE_DIR / ".env")
+
+
+def cosine_similarity(a: np.ndarray, b: np.ndarray) -> float:
+    """计算两个向量的余弦相似度。"""
+    norm_a = np.linalg.norm(a)
+    norm_b = np.linalg.norm(b)
+    if norm_a == 0 or norm_b == 0:
+        return 0.0
+    return float(np.dot(a, b) / (norm_a * norm_b))
+
+
+def vector_search(conn: sqlite3.Connection, query_embedding: list[float], top_k: int = 5) -> list[dict]:
+    """向量检索：在 embeddings 表中找到与查询向量最相似的文章。"""
+    query_vec = np.array(query_embedding, dtype=np.float32)
+    rows = conn.execute(
+        """
+        SELECT e.article_id, e.embedding, e.chunk_text,
+               a.title, a.source_name, a.source_type, a.published_at, a.summary,
+               x.result_json
+        FROM embeddings e
+        JOIN articles a ON a.id = e.article_id
+        LEFT JOIN ai_analyses x ON x.id = (
+            SELECT x2.id FROM ai_analyses x2
+            WHERE x2.article_id = e.article_id AND x2.analysis_type = 'editorial'
+            ORDER BY x2.updated_at DESC, x2.id DESC LIMIT 1
+        )
+        WHERE a.crawl_status = 'success'
+        """
+    ).fetchall()
+
+    scored = []
+    for row in rows:
+        emb = json.loads(row["embedding"])
+        emb_vec = np.array(emb, dtype=np.float32)
+        score = cosine_similarity(query_vec, emb_vec)
+        scored.append((score, row))
+    scored.sort(key=lambda x: x[0], reverse=True)
+
+    results = []
+    seen_articles = set()
+    for score, row in scored:
+        if row["article_id"] in seen_articles:
+            continue
+        seen_articles.add(row["article_id"])
+        analysis = json.loads(row["result_json"]) if row["result_json"] else {}
+        results.append({
+            "article_id": row["article_id"],
+            "title": row["title"],
+            "source": row["source_name"] or row["source_type"],
+            "date": (row["published_at"] or "")[:10],
+            "summary": analysis.get("summary") or row["summary"] or "",
+            "key_points": analysis.get("key_points", []),
+            "content": row["chunk_text"] or "",
+            "score": score,
+        })
+        if len(results) >= top_k:
+            break
+    return results
+
+
+def generate_answer(question: str, contexts: list[dict]) -> dict:
+    """基于检索到的上下文，调用千问生成回答。"""
+    api_key = get_api_key()
+    model = os.getenv("QWEN_MODEL", "qwen-plus")
+
+    # 构建上下文
+    context_text = ""
+    sources = []
+    for i, ctx in enumerate(contexts, start=1):
+        context_text += f"\n--- 来源{i} ---\n标题：{ctx['title']}\n摘要：{ctx['summary']}\n要点：{'; '.join(ctx['key_points'][:3])}\n原文片段：{ctx['content']}\n"
+        sources.append({
+            "title": ctx["title"],
+            "source": ctx["source"],
+            "date": ctx["date"],
+        })
+
+    prompt = f"""你是餐饮行业知识库的问答助手。根据以下知识库中的情报内容回答用户问题。
+
+要求：
+1. 回答必须基于提供的来源内容，不要编造信息
+2. 如果来源内容不足以回答问题，明确说明"知识库中暂无直接相关信息"
+3. 回答要结构化、简洁、有经营参考价值
+4. 用中文回答
+
+知识库来源：
+{context_text}
+
+用户问题：{question}"""
+
+    payload = json.dumps({
+        "model": model,
+        "messages": [
+            {"role": "system", "content": "你是餐饮行业知识库问答助手，基于库内情报给出有出处的专业回答。"},
+            {"role": "user", "content": prompt},
+        ],
+        "temperature": 0.3,
+    }, ensure_ascii=False).encode("utf-8")
+
+    req = urllib.request.Request(
+        DASHSCOPE_CHAT_URL,
+        data=payload,
+        headers={"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"},
+        method="POST",
+    )
+    try:
+        with urllib.request.urlopen(req, timeout=90) as response:
+            body = json.loads(response.read().decode("utf-8"))
+    except urllib.error.HTTPError as exc:
+        detail = exc.read().decode("utf-8", errors="replace")[:500]
+        raise RuntimeError(f"千问接口返回 {exc.code}：{detail}") from exc
+
+    answer = body["choices"][0]["message"]["content"].strip()
+    return {
+        "answer": answer,
+        "sources": sources,
+        "contexts": contexts,
+    }
+
+
+def generate_answer_stream(question: str, contexts: list[dict]):
+    """流式生成回答，逐 token yield。"""
+    api_key = get_api_key()
+    model = os.getenv("QWEN_MODEL", "qwen-plus")
+
+    context_text = ""
+    sources = []
+    for i, ctx in enumerate(contexts, start=1):
+        context_text += f"\n--- 来源{i} ---\n标题：{ctx['title']}\n摘要：{ctx['summary']}\n要点：{'; '.join(ctx['key_points'][:3])}\n原文片段：{ctx['content']}\n"
+        sources.append({
+            "title": ctx["title"],
+            "source": ctx["source"],
+            "date": ctx["date"],
+        })
+
+    prompt = f"""你是餐饮行业知识库的问答助手。根据以下知识库中的情报内容回答用户问题。
+
+要求：
+1. 回答必须基于提供的来源内容，不要编造信息
+2. 如果来源内容不足以回答问题，明确说明"知识库中暂无直接相关信息"
+3. 回答要结构化、简洁、有经营参考价值
+4. 用中文回答，使用 Markdown 格式
+
+知识库来源：
+{context_text}
+
+用户问题：{question}"""
+
+    payload = json.dumps({
+        "model": model,
+        "messages": [
+            {"role": "system", "content": "你是餐饮行业知识库问答助手，基于库内情报给出有出处的专业回答。使用 Markdown 格式输出。"},
+            {"role": "user", "content": prompt},
+        ],
+        "temperature": 0.3,
+        "stream": True,
+        "stream_options": {"include_usage": False},
+    }, ensure_ascii=False).encode("utf-8")
+
+    req = urllib.request.Request(
+        DASHSCOPE_CHAT_URL,
+        data=payload,
+        headers={
+            "Authorization": f"Bearer {api_key}",
+            "Content-Type": "application/json",
+            "Accept": "text/event-stream",
+        },
+        method="POST",
+    )
+    try:
+        resp = urllib.request.urlopen(req, timeout=120)
+    except urllib.error.HTTPError as exc:
+        detail = exc.read().decode("utf-8", errors="replace")[:500]
+        raise RuntimeError(f"千问接口返回 {exc.code}：{detail}") from exc
+
+    buffer = b""
+    while True:
+        chunk = resp.read(4096)
+        if not chunk:
+            break
+        buffer += chunk
+        while b"\n" in buffer:
+            line, buffer = buffer.split(b"\n", 1)
+            line = line.strip()
+            if not line or line == b"data: [DONE]":
+                continue
+            if line.startswith(b"data:"):
+                try:
+                    data = json.loads(line[5:].strip())
+                    delta = data.get("choices", [{}])[0].get("delta", {})
+                    content = delta.get("content")
+                    if content:
+                        yield content
+                except (json.JSONDecodeError, IndexError, KeyError):
+                    pass
+    resp.close()
+
+    yield json.dumps({"__sources__": sources}, ensure_ascii=False)
+
+
+def rag_qa(question: str, top_k: int = 5) -> dict:
+    """完整的 RAG 问答流程：向量检索 → LLM 生成。"""
+    db_path = Path(os.getenv("CIBANK_DB", str(BASE_DIR / "data" / "cibank.db")))
+    conn = sqlite3.connect(db_path)
+    conn.row_factory = sqlite3.Row
+
+    # 检查是否有向量数据
+    count = conn.execute("SELECT COUNT(*) FROM embeddings").fetchone()[0]
+    if count == 0:
+        # 无向量数据时使用关键词检索作为降级方案
+        return keyword_fallback(conn, question)
+
+    # 1. 将问题转为向量
+    query_embeddings = get_embeddings([question])
+
+    # 2. 向量检索
+    contexts = vector_search(conn, query_embeddings[0], top_k)
+
+    if not contexts:
+        conn.close()
+        return {
+            "answer": "知识库中暂无与您问题相关的内容。请先运行采集和分析流程入库更多情报。",
+            "sources": [],
+            "contexts": [],
+        }
+
+    # 3. LLM 生成回答
+    result = generate_answer(question, contexts)
+    conn.close()
+    return result
+
+
+def rag_qa_stream(question: str, top_k: int = 5):
+    """流式 RAG 问答：先检索，再流式生成。yield (type, data) 元组。"""
+    db_path = Path(os.getenv("CIBANK_DB", str(BASE_DIR / "data" / "cibank.db")))
+    conn = sqlite3.connect(db_path)
+    conn.row_factory = sqlite3.Row
+
+    count = conn.execute("SELECT COUNT(*) FROM embeddings").fetchone()[0]
+    if count == 0:
+        conn.close()
+        yield ("error", "知识库中暂无向量数据，请先运行采集和分析流程。")
+        return
+
+    query_embeddings = get_embeddings([question])
+    contexts = vector_search(conn, query_embeddings[0], top_k)
+    conn.close()
+
+    if not contexts:
+        yield ("error", "知识库中暂无与您问题相关的内容。")
+        return
+
+    sources = [
+        {"title": ctx["title"], "source": ctx["source"], "date": ctx["date"]}
+        for ctx in contexts
+    ]
+    yield ("sources", sources)
+
+    for token in generate_answer_stream(question, contexts):
+        yield ("token", token)
+
+
+def keyword_fallback(conn: sqlite3.Connection, question: str) -> dict:
+    """无向量数据时的关键词检索降级方案。"""
+    rows = conn.execute(
+        """
+        SELECT a.id, a.title, a.source_name, a.source_type, a.published_at, a.summary,
+               a.content, x.result_json
+        FROM articles a
+        LEFT JOIN ai_analyses x ON x.article_id = a.id AND x.analysis_type = 'editorial'
+        WHERE a.crawl_status = 'success'
+          AND (a.title LIKE ? OR a.summary LIKE ? OR a.content LIKE ?)
+        ORDER BY a.published_at DESC
+        LIMIT 5
+        """,
+        (f"%{question}%", f"%{question}%", f"%{question}%"),
+    ).fetchall()
+
+    if not rows:
+        conn.close()
+        return {
+            "answer": "知识库中暂无与您问题相关的内容。请先运行采集和分析流程入库更多情报。",
+            "sources": [],
+            "contexts": [],
+        }
+
+    contexts = []
+    for row in rows:
+        analysis = json.loads(row["result_json"]) if row["result_json"] else {}
+        contexts.append({
+            "article_id": row["id"],
+            "title": row["title"],
+            "source": row["source_name"] or row["source_type"],
+            "date": (row["published_at"] or "")[:10],
+            "summary": analysis.get("summary") or row["summary"] or "",
+            "key_points": analysis.get("key_points", []),
+            "content": row["content"][:3000] if row["content"] else "",
+            "score": 0.5,
+        })
+    conn.close()
+    return generate_answer(question, contexts)
+
+
+if __name__ == "__main__":
+    import sys
+    question = " ".join(sys.argv[1:]) or "现在开茶饮店还有机会吗？"
+    result = rag_qa(question)
+    print(json.dumps(result, ensure_ascii=False, indent=2))