241 lines
10 KiB
Python
241 lines
10 KiB
Python
"""Implement infrastructure support for openai compatible answer generator."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import time
|
|
from typing import Generator
|
|
|
|
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.
|
|
|
|
|
|
# Fallback system prompts used when no rich template matches.
|
|
_FALLBACK_PROMPTS = {
|
|
"default": "你是法规知识问答助手。请仅依据提供的上下文回答；如果上下文不足，明确说明。",
|
|
"compliance_qa": "你是法规合规问答助手。优先引用给定法规原文，回答要准确、克制，并注明依据来源。",
|
|
}
|
|
|
|
|
|
class OpenAICompatibleAnswerGenerator(AnswerGenerator):
|
|
"""Represent the Open A I Compatible Answer Generator type."""
|
|
|
|
@staticmethod
|
|
def _estimate_tokens(text: str) -> int:
|
|
"""Estimate token count for mixed Chinese/English text.
|
|
|
|
Chinese chars are ~1.5 chars/token; ASCII is ~4 chars/token.
|
|
"""
|
|
chinese = sum(1 for c in text if "一" <= c <= "鿿")
|
|
other = len(text) - chinese
|
|
return int(chinese / 1.5 + other / 4) + 1
|
|
|
|
def _build_messages(
|
|
self,
|
|
*,
|
|
query: str,
|
|
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]:
|
|
"""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"章节: {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 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"]})
|
|
|
|
# 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:
|
|
"""Return whether the prompt context had to omit retrieved chunks to fit the token budget."""
|
|
if not retrieved_chunks:
|
|
return False
|
|
estimated_total_tokens = sum(
|
|
self._estimate_tokens(
|
|
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}"
|
|
)
|
|
for idx, chunk in enumerate(retrieved_chunks, start=1)
|
|
)
|
|
return estimated_total_tokens > context_tokens
|
|
|
|
def _sources(self, chunks: list[RetrievedChunk]) -> list[AnswerSource]:
|
|
"""Handle sources for this module for the Open A I Compatible Answer Generator instance."""
|
|
return [
|
|
AnswerSource(
|
|
doc_id=chunk.doc_id,
|
|
doc_title=chunk.doc_title,
|
|
chunk_id=chunk.chunk_id,
|
|
chunk_type=chunk.chunk_type,
|
|
section_title=chunk.section_title,
|
|
page_start=chunk.page_start,
|
|
page_end=chunk.page_end,
|
|
section_level=chunk.section_level,
|
|
chunk_index=chunk.chunk_index,
|
|
piece_index=chunk.piece_index,
|
|
score=chunk.score,
|
|
text=chunk.text,
|
|
metadata=chunk.metadata,
|
|
)
|
|
for chunk in chunks
|
|
]
|
|
|
|
def generate(
|
|
self,
|
|
*,
|
|
query: str,
|
|
retrieved_chunks: list[RetrievedChunk],
|
|
history: list[dict[str, str]] | None = None,
|
|
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()
|
|
messages, context_tokens = self._build_messages(
|
|
query=query,
|
|
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)
|
|
latency_ms = int((time.time() - start) * 1000)
|
|
return AnswerResult(
|
|
answer=response.content if response.is_success else "",
|
|
sources=self._sources(retrieved_chunks),
|
|
model=response.model or (model or settings.llm_model),
|
|
latency_ms=latency_ms,
|
|
retrieved_count=len(retrieved_chunks),
|
|
context_tokens=context_tokens,
|
|
truncated=self._is_context_truncated(
|
|
retrieved_chunks=retrieved_chunks,
|
|
context_tokens=context_tokens,
|
|
),
|
|
error=response.error,
|
|
)
|
|
|
|
def stream_generate(
|
|
self,
|
|
*,
|
|
query: str,
|
|
retrieved_chunks: list[RetrievedChunk],
|
|
history: list[dict[str, str]] | None = None,
|
|
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()
|
|
messages, context_tokens = self._build_messages(
|
|
query=query,
|
|
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}
|
|
client = get_llm_client(provider=provider or settings.llm_provider, model=model or settings.llm_model)
|
|
answer_parts: list[str] = []
|
|
try:
|
|
if hasattr(client, "stream_chat"):
|
|
for chunk in client.stream_chat(messages):
|
|
answer_parts.append(chunk)
|
|
yield {"event": "content", "data": chunk}
|
|
else:
|
|
response = client.chat(messages)
|
|
answer_parts.append(response.content)
|
|
yield {"event": "content", "data": response.content}
|
|
except Exception as exc:
|
|
yield {"event": "error", "data": str(exc)}
|
|
return
|
|
yield {
|
|
"event": "done",
|
|
"data": {
|
|
"latency_ms": int((time.time() - start) * 1000),
|
|
"retrieved_count": len(retrieved_chunks),
|
|
"context_tokens": context_tokens,
|
|
"model": model or settings.llm_model,
|
|
},
|
|
}
|