AIRegulation/AIRegulation-DocAnalysis
7
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add LLM token

Browse Source
This commit is contained in:
wangwei
2026-07-02 22:03:39 +08:00
parent e3afb8a07a
commit 52e67b0e7b
36 changed files with 2392 additions and 394 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/app/infrastructure/llm/openai_compatible_answer_generator.py
+68 -11
View File
@@ -9,10 +9,12 @@ from app.config.settings import settings
from app.domain.conversation import AnswerGenerator, AnswerResult, AnswerSource
from app.domain.retrieval import RetrievedChunk
from app.services.llm.llm_factory import get_llm_client
from app.services.rag.prompt_templates import PromptTemplates
# Keep adapter behavior explicit so integration details remain easy to audit.
PROMPT_TEMPLATES = {
# Fallback system prompts used when no rich template matches.
_FALLBACK_PROMPTS = {
"default": "你是法规知识问答助手。请仅依据提供的上下文回答；如果上下文不足，明确说明。",
"compliance_qa": "你是法规合规问答助手。优先引用给定法规原文，回答要准确、克制，并注明依据来源。",
}
@@ -38,33 +40,80 @@ class OpenAICompatibleAnswerGenerator(AnswerGenerator):
retrieved_chunks: list[RetrievedChunk],
history: list[dict[str, str]] | None,
prompt_template: str | None,
context_text: str | None = None,
context_filename: str | None = None,
) -> tuple[list[dict[str, str]], int]:
"""Handle build messages for this module for the Open A I Compatible Answer Generator instance."""
system_prompt = PROMPT_TEMPLATES.get(prompt_template or "compliance_qa", PROMPT_TEMPLATES["default"])
"""Build the message list to send to the LLM.
When context_text is provided the user's document is injected as a
dedicated section BEFORE the retrieved regulation chunks so the LLM
can reason about the document directly while still referencing regulations.
The retrieval step uses only the user's question, not the document text,
so embedding quality is preserved.
System prompt selection priority:
1. Rich template from PromptTemplates (compliance_qa / comparison /
compliance_check / clause_interpretation / …)
2. Fallback hardcoded prompt when no rich template matches.
"""
# Look up the rich template first; fall back to simple hardcoded prompts.
tpl_name = prompt_template or "compliance_qa"
rich_tpl = PromptTemplates.get_template(tpl_name)
if rich_tpl:
system_prompt = rich_tpl.system_prompt
else:
system_prompt = _FALLBACK_PROMPTS.get(tpl_name, _FALLBACK_PROMPTS["default"])
context_blocks = []
context_tokens = 0
# ── User document context (if attached) ───────────────────────────────
if context_text and context_text.strip():
doc_label = f"附件文档：{context_filename}" if context_filename else "附件文档"
doc_block = f"[{doc_label}]\n{context_text.strip()}"
doc_tokens = self._estimate_tokens(doc_block)
# Reserve at most half the context budget for the user document
half_budget = settings.rag_max_context_tokens // 2
if doc_tokens > half_budget:
# Truncate document to fit half the budget
ratio = half_budget / doc_tokens
doc_block = doc_block[: int(len(doc_block) * ratio)] + "\n…（文档已截断）"
doc_tokens = half_budget
context_blocks.append(doc_block)
context_tokens += doc_tokens
# ── Retrieved regulation chunks ────────────────────────────────────────
remaining_budget = settings.rag_max_context_tokens - context_tokens
for idx, chunk in enumerate(retrieved_chunks, start=1):
block = (
f"[{idx}] 文档: {chunk.doc_title}\n"
f"[法规{idx}] 文档: {chunk.doc_title}\n"
f"章节: {chunk.section_title or '未标注'}\n"
f"页码: {chunk.page_start}" + (f"-{chunk.page_end}" if chunk.page_end and chunk.page_end != chunk.page_start else "") + "\n"
f"内容: {chunk.text}"
)
block_tokens = self._estimate_tokens(block)
if context_tokens + block_tokens > settings.rag_max_context_tokens:
if block_tokens > remaining_budget:
break
remaining_budget -= block_tokens
context_tokens += block_tokens
context_blocks.append(block)
context = "\n\n".join(context_blocks)
messages = [{"role": "system", "content": system_prompt}]
for item in history or []:
messages.append({"role": item["role"], "content": item["content"]})
messages.append(
{
"role": "user",
"content": f"问题：{query}\n\n参考上下文：\n{context}\n\n请在回答后给出简要引用编号。",
}
)
# Craft the user turn differently when a document is attached
if context_text and context_text.strip():
user_content = (
f"问题：{query}\n\n"
f"请先基于上方附件文档内容进行分析，再结合法规参考上下文给出合规评估。"
f"\n\n参考上下文：\n{context}\n\n"
f"请在回答中注明引用来源编号（如适用）。"
)
else:
user_content = f"问题：{query}\n\n参考上下文：\n{context}\n\n请在回答后给出简要引用编号。"
messages.append({"role": "user", "content": user_content})
return messages, context_tokens
def _is_context_truncated(self, *, retrieved_chunks: list[RetrievedChunk], context_tokens: int) -> bool:
@@ -112,6 +161,8 @@ class OpenAICompatibleAnswerGenerator(AnswerGenerator):
provider: str | None = None,
model: str | None = None,
prompt_template: str | None = None,
context_text: str | None = None,
context_filename: str | None = None,
) -> AnswerResult:
"""Handle generate for the Open A I Compatible Answer Generator instance."""
start = time.time()
@@ -120,6 +171,8 @@ class OpenAICompatibleAnswerGenerator(AnswerGenerator):
retrieved_chunks=retrieved_chunks,
history=history,
prompt_template=prompt_template,
context_text=context_text,
context_filename=context_filename,
)
client = get_llm_client(provider=provider or settings.llm_provider, model=model or settings.llm_model)
response = client.chat(messages)
@@ -147,6 +200,8 @@ class OpenAICompatibleAnswerGenerator(AnswerGenerator):
provider: str | None = None,
model: str | None = None,
prompt_template: str | None = None,
context_text: str | None = None,
context_filename: str | None = None,
) -> Generator[dict, None, AnswerResult]:
"""Stream generate for the Open A I Compatible Answer Generator instance."""
start = time.time()
@@ -155,6 +210,8 @@ class OpenAICompatibleAnswerGenerator(AnswerGenerator):
retrieved_chunks=retrieved_chunks,
history=history,
prompt_template=prompt_template,
context_text=context_text,
context_filename=context_filename,
)
sources = [source.__dict__ for source in self._sources(retrieved_chunks)]
yield {"event": "sources", "data": sources}