86 lines
3.5 KiB
Python
86 lines
3.5 KiB
Python
import logging
|
||
|
||
import httpx
|
||
|
||
from app.config import LOW_SCORE, llm_config
|
||
|
||
logger = logging.getLogger("learn_rag")
|
||
|
||
SYSTEM_PROMPT = """你是「知迹」的知识库问答助手。只能根据给出的知识切片回答，不要补充切片之外的信息，也不要编造。
|
||
如果切片不足以回答，就说明知识库里没有这项信息。
|
||
用自然的中文直接回答问题，不要使用 Markdown，不要写标题。
|
||
引用事实时在相关句末标注编号，格式必须是〔来源1〕这种全角括号。编号只能使用给定切片的编号。"""
|
||
|
||
|
||
def build_context(chunks: list[dict]) -> str:
|
||
blocks = []
|
||
for chunk in chunks:
|
||
blocks.append(
|
||
f"〔来源{chunk['n']}〕《{chunk['title']}》相关度 {chunk['score']:.0%}\n{chunk['text']}"
|
||
)
|
||
return "\n\n".join(blocks)
|
||
|
||
|
||
def extractive_answer(chunks: list[dict], low_confidence: bool) -> str:
|
||
if not chunks:
|
||
return "知识库里还没有可以检索的内容。"
|
||
lead = ""
|
||
if low_confidence:
|
||
lead = "最接近的切片和问题只是部分相关，下面按原文摘出，请以右侧切片为准。\n\n"
|
||
lines = [lead + chunks[0]["text"].strip() + f"〔来源{chunks[0]['n']}〕"]
|
||
for chunk in chunks[1:2]:
|
||
if chunks[0]["score"] - chunk["score"] <= 0.08:
|
||
lines.append("\n" + chunk["text"].strip() + f"〔来源{chunk['n']}〕")
|
||
return "\n".join(lines).strip()
|
||
|
||
|
||
def generate(question: str, chunks: list[dict]) -> tuple[str, str, str | None, bool]:
|
||
low_confidence = not chunks or chunks[0]["score"] < LOW_SCORE
|
||
if not chunks:
|
||
return "知识库里还没有可以检索的内容。", "extractive", None, True
|
||
|
||
settings = llm_config()
|
||
if settings is None:
|
||
notice = "还没有配置大模型，回答直接摘自相关度最高的切片。在项目目录的 .env 中填入 OPENAI_API_KEY 后重启，即可生成连贯回答。"
|
||
return extractive_answer(chunks, low_confidence), "extractive", notice, low_confidence
|
||
|
||
caution = ""
|
||
if low_confidence:
|
||
caution = "注意：最高相关度偏低。如果切片不能直接支持答案，请明确说知识库里没有这项信息。\n\n"
|
||
user_prompt = f"{caution}知识切片如下：\n\n{build_context(chunks)}\n\n问题：{question}"
|
||
try:
|
||
answer = call_llm(settings, user_prompt)
|
||
except Exception:
|
||
logger.exception("大模型调用失败")
|
||
notice = "大模型暂时不可用，已改为直接摘录检索到的切片。"
|
||
return extractive_answer(chunks, low_confidence), "extractive", notice, low_confidence
|
||
return answer, "llm", None, low_confidence
|
||
|
||
|
||
def call_llm(settings: dict, user_prompt: str) -> str:
|
||
response = httpx.post(
|
||
f"{settings['base']}/chat/completions",
|
||
headers={"Authorization": f"Bearer {settings['key']}"},
|
||
json={
|
||
"model": settings["model"],
|
||
"temperature": 0.2,
|
||
"max_tokens": 800,
|
||
"messages": [
|
||
{"role": "system", "content": SYSTEM_PROMPT},
|
||
{"role": "user", "content": user_prompt},
|
||
],
|
||
},
|
||
timeout=httpx.Timeout(60.0, connect=10.0),
|
||
)
|
||
response.raise_for_status()
|
||
message = response.json()["choices"][0]["message"]
|
||
content = message.get("content")
|
||
if isinstance(content, list):
|
||
content = "".join(
|
||
part.get("text", "") if isinstance(part, dict) else str(part) for part in content
|
||
)
|
||
text = (content or "").strip()
|
||
if not text:
|
||
raise ValueError("模型返回了空回答")
|
||
return text
|