Add V2 RAG answer skeleton
This commit is contained in:
@@ -72,7 +72,7 @@ def completions(
|
||||
db: Session = Depends(get_db),
|
||||
current_user: User = Depends(get_current_user),
|
||||
) -> StreamingResponse:
|
||||
answer = ChatService.create_mock_answer(db, current_user, payload.sessionId, payload.message)
|
||||
answer = ChatService.create_answer(db, current_user, payload.sessionId, payload.message)
|
||||
return StreamingResponse(_sse_chunks(answer), media_type="text/event-stream")
|
||||
|
||||
|
||||
|
||||
@@ -32,6 +32,7 @@ class Settings(BaseSettings):
|
||||
mock_sms_enabled: bool = True
|
||||
mock_sms_code: str = "123456"
|
||||
sms_code_expire_minutes: int = 5
|
||||
mock_rag_enabled: bool = True
|
||||
|
||||
default_daily_chat_limit: int = 100
|
||||
default_user_name_prefix: str = "用户"
|
||||
|
||||
@@ -0,0 +1,39 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from app.models.logs import AiRequestLog
|
||||
|
||||
|
||||
class AiRequestLogService:
|
||||
@staticmethod
|
||||
def write_success(
|
||||
db: Session,
|
||||
*,
|
||||
session_id: int,
|
||||
message_id: int | None,
|
||||
user_id: int,
|
||||
model_name: str,
|
||||
prompt: str,
|
||||
knowledge_ids: str,
|
||||
retrieve_count: int,
|
||||
input_token: int,
|
||||
output_token: int,
|
||||
cost_ms: int,
|
||||
) -> None:
|
||||
db.add(
|
||||
AiRequestLog(
|
||||
session_id=session_id,
|
||||
message_id=message_id,
|
||||
user_id=user_id,
|
||||
model_name=model_name,
|
||||
prompt=prompt,
|
||||
knowledge_ids=knowledge_ids,
|
||||
retrieve_count=retrieve_count,
|
||||
input_token=input_token,
|
||||
output_token=output_token,
|
||||
total_token=input_token + output_token,
|
||||
cost_ms=cost_ms,
|
||||
status="SUCCESS",
|
||||
)
|
||||
)
|
||||
@@ -1,6 +1,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import UTC, datetime
|
||||
from time import perf_counter
|
||||
|
||||
from fastapi import HTTPException, status
|
||||
from sqlalchemy import select
|
||||
@@ -8,6 +9,9 @@ from sqlalchemy.orm import Session
|
||||
|
||||
from app.models.chat import ChatMessage, ChatSession
|
||||
from app.models.user import User
|
||||
from app.services.ai_request_log_service import AiRequestLogService
|
||||
from app.services.model_service import ModelClientService
|
||||
from app.services.rag_service import RagService
|
||||
|
||||
|
||||
class ChatService:
|
||||
@@ -64,36 +68,65 @@ class ChatService:
|
||||
db.commit()
|
||||
|
||||
@staticmethod
|
||||
def create_mock_answer(db: Session, user: User, session_id: int, question: str) -> str:
|
||||
def create_answer(db: Session, user: User, session_id: int, question: str) -> str:
|
||||
session = ChatService._get_user_session(db, user, session_id)
|
||||
ChatService._ensure_quota(user)
|
||||
|
||||
now = _now()
|
||||
answer = (
|
||||
"这是 mock AI 回复。当前阶段已经跑通后端聊天链路，后续会替换为：权限过滤、飞书实时检索、"
|
||||
"Prompt 组装和真实大模型 SSE 输出。"
|
||||
)
|
||||
normalized_question = question.strip()
|
||||
user_message = ChatMessage(
|
||||
session_id=session.id,
|
||||
user_id=user.id,
|
||||
role="user",
|
||||
content=question.strip(),
|
||||
content=normalized_question,
|
||||
message_status="FINISHED",
|
||||
created_at=now,
|
||||
)
|
||||
db.add(user_message)
|
||||
db.flush()
|
||||
|
||||
started_at = perf_counter()
|
||||
rag_result = RagService.build_result(db, user, normalized_question)
|
||||
completion = ModelClientService.complete(db, rag_result)
|
||||
cost_ms = int((perf_counter() - started_at) * 1000)
|
||||
|
||||
assistant_message = ChatMessage(
|
||||
session_id=session.id,
|
||||
user_id=user.id,
|
||||
role="assistant",
|
||||
content=answer,
|
||||
content=completion.answer,
|
||||
message_status="FINISHED",
|
||||
token_input=completion.input_token,
|
||||
token_output=completion.output_token,
|
||||
response_time_ms=cost_ms,
|
||||
model_id=completion.model_id,
|
||||
created_at=now,
|
||||
)
|
||||
db.add(assistant_message)
|
||||
db.flush()
|
||||
|
||||
session.message_count += 2
|
||||
session.last_message_at = now
|
||||
if session.title == "新聊天":
|
||||
session.title = _title_from_question(question)
|
||||
db.add_all([user_message, assistant_message, session])
|
||||
session.title = _title_from_question(normalized_question)
|
||||
|
||||
user.daily_chat_used += 1
|
||||
db.add_all([session, user])
|
||||
AiRequestLogService.write_success(
|
||||
db,
|
||||
session_id=session.id,
|
||||
message_id=assistant_message.id,
|
||||
user_id=user.id,
|
||||
model_name=completion.model_name,
|
||||
prompt=rag_result.prompt,
|
||||
knowledge_ids=rag_result.knowledge_ids,
|
||||
retrieve_count=len(rag_result.chunks),
|
||||
input_token=completion.input_token,
|
||||
output_token=completion.output_token,
|
||||
cost_ms=cost_ms,
|
||||
)
|
||||
db.commit()
|
||||
return answer
|
||||
return completion.answer
|
||||
|
||||
@staticmethod
|
||||
def stop_generation(db: Session, user: User, session_id: int) -> None:
|
||||
@@ -106,6 +139,11 @@ class ChatService:
|
||||
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="会话不存在")
|
||||
return session
|
||||
|
||||
@staticmethod
|
||||
def _ensure_quota(user: User) -> None:
|
||||
if user.daily_chat_used >= user.daily_chat_limit:
|
||||
raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="今日提问次数已用完")
|
||||
|
||||
|
||||
def _now() -> datetime:
|
||||
return datetime.now(UTC).replace(tzinfo=None)
|
||||
|
||||
@@ -0,0 +1,60 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from datetime import UTC, datetime
|
||||
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from app.core.config import get_settings
|
||||
from app.models.knowledge import Knowledge, UserKnowledgePermission
|
||||
from app.models.user import User
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class KnowledgeScope:
|
||||
id: int
|
||||
name: str
|
||||
feishu_space_id: str
|
||||
feishu_node_id: str
|
||||
is_mock: bool = False
|
||||
|
||||
|
||||
class KnowledgeAccessService:
|
||||
@staticmethod
|
||||
def get_allowed_knowledge(db: Session, user: User) -> list[KnowledgeScope]:
|
||||
now = datetime.now(UTC).replace(tzinfo=None)
|
||||
rows = db.execute(
|
||||
select(Knowledge)
|
||||
.join(UserKnowledgePermission, UserKnowledgePermission.knowledge_id == Knowledge.id)
|
||||
.where(
|
||||
UserKnowledgePermission.user_id == user.id,
|
||||
Knowledge.status == 1,
|
||||
(UserKnowledgePermission.effective_at.is_(None))
|
||||
| (UserKnowledgePermission.effective_at <= now),
|
||||
(UserKnowledgePermission.expired_at.is_(None))
|
||||
| (UserKnowledgePermission.expired_at >= now),
|
||||
)
|
||||
.order_by(Knowledge.id.asc())
|
||||
).scalars()
|
||||
scopes = [
|
||||
KnowledgeScope(
|
||||
id=knowledge.id,
|
||||
name=knowledge.name,
|
||||
feishu_space_id=knowledge.feishu_space_id,
|
||||
feishu_node_id=knowledge.feishu_node_id,
|
||||
)
|
||||
for knowledge in rows
|
||||
]
|
||||
if scopes or not get_settings().mock_rag_enabled:
|
||||
return scopes
|
||||
|
||||
return [
|
||||
KnowledgeScope(
|
||||
id=0,
|
||||
name="开发阶段默认知识库",
|
||||
feishu_space_id="mock-space",
|
||||
feishu_node_id="mock-node",
|
||||
is_mock=True,
|
||||
)
|
||||
]
|
||||
@@ -0,0 +1,57 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from app.models.ai_config import ModelConfig
|
||||
from app.services.rag_service import NO_HIT_ANSWER, RagResult
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ModelCompletion:
|
||||
answer: str
|
||||
model_id: int | None
|
||||
model_name: str
|
||||
input_token: int
|
||||
output_token: int
|
||||
|
||||
|
||||
class ModelClientService:
|
||||
@staticmethod
|
||||
def complete(db: Session, rag_result: RagResult) -> ModelCompletion:
|
||||
model = db.scalar(
|
||||
select(ModelConfig)
|
||||
.where(ModelConfig.enabled == 1)
|
||||
.order_by(ModelConfig.id.desc())
|
||||
.limit(1)
|
||||
)
|
||||
model_name = model.model_name if model is not None else "mock-model"
|
||||
answer = _mock_answer(rag_result)
|
||||
return ModelCompletion(
|
||||
answer=answer,
|
||||
model_id=model.id if model is not None else None,
|
||||
model_name=model_name,
|
||||
input_token=_rough_token_count(rag_result.prompt),
|
||||
output_token=_rough_token_count(answer),
|
||||
)
|
||||
|
||||
|
||||
def _mock_answer(rag_result: RagResult) -> str:
|
||||
if not rag_result.is_hit:
|
||||
return NO_HIT_ANSWER
|
||||
|
||||
summaries = "\n".join(
|
||||
f"- {chunk.content}\n 来源：{chunk.knowledge_name} / {chunk.title}" for chunk in rag_result.chunks
|
||||
)
|
||||
return (
|
||||
"根据当前已授权知识库，整理到的信息如下：\n\n"
|
||||
f"{summaries}\n\n"
|
||||
"当前仍是 mock RAG + mock 模型阶段，后续会把检索服务替换为飞书实时检索，"
|
||||
"把模型服务替换为后台启用的大模型配置。"
|
||||
)
|
||||
|
||||
|
||||
def _rough_token_count(text: str) -> int:
|
||||
return max(1, len(text.strip()) // 2)
|
||||
124
ai_knowledge_base_v2/apps/backend/app/services/rag_service.py
Normal file
124
ai_knowledge_base_v2/apps/backend/app/services/rag_service.py
Normal file
@@ -0,0 +1,124 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from app.models.ai_config import Prompt
|
||||
from app.models.user import User
|
||||
from app.services.knowledge_service import KnowledgeAccessService, KnowledgeScope
|
||||
|
||||
NO_HIT_ANSWER = "当前知识库中未检索到相关内容，请联系管理员补充相关知识。"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class RetrievedChunk:
|
||||
knowledge_id: int
|
||||
knowledge_name: str
|
||||
title: str
|
||||
content: str
|
||||
source_url: str | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class RagResult:
|
||||
question: str
|
||||
knowledge_scopes: list[KnowledgeScope]
|
||||
chunks: list[RetrievedChunk]
|
||||
prompt: str
|
||||
|
||||
@property
|
||||
def is_hit(self) -> bool:
|
||||
return len(self.chunks) > 0
|
||||
|
||||
@property
|
||||
def knowledge_ids(self) -> str:
|
||||
return ",".join(str(scope.id) for scope in self.knowledge_scopes)
|
||||
|
||||
|
||||
class RagService:
|
||||
@staticmethod
|
||||
def build_result(db: Session, user: User, question: str) -> RagResult:
|
||||
scopes = KnowledgeAccessService.get_allowed_knowledge(db, user)
|
||||
chunks = FeishuKnowledgeService.retrieve(question, scopes)
|
||||
prompt = PromptService.build_prompt(db, question, chunks)
|
||||
return RagResult(question=question, knowledge_scopes=scopes, chunks=chunks, prompt=prompt)
|
||||
|
||||
|
||||
class FeishuKnowledgeService:
|
||||
_mock_documents = [
|
||||
{
|
||||
"title": "一期产品目标",
|
||||
"keywords": {"一期", "目标", "问答", "登录", "知识库", "飞书", "用户端", "后台"},
|
||||
"content": "一期要交付企业飞书知识库 AI 问答系统，核心包括用户登录、AI 问答、权限内知识库回答和后台管理能力。",
|
||||
},
|
||||
{
|
||||
"title": "权限过滤规则",
|
||||
"keywords": {"权限", "授权", "越权", "用户", "知识库", "过期", "禁用"},
|
||||
"content": "RAG 检索前必须先过滤用户授权知识库，只允许未禁用、未过期、当前用户有权限的知识库参与回答。",
|
||||
},
|
||||
{
|
||||
"title": "无命中兜底规则",
|
||||
"keywords": {"无命中", "兜底", "编造", "检索", "答案", "命中"},
|
||||
"content": f"当知识库没有命中相关内容时，系统必须固定返回：{NO_HIT_ANSWER}",
|
||||
},
|
||||
{
|
||||
"title": "技术实现边界",
|
||||
"keywords": {"模型", "prompt", "大模型", "日志", "sse", "流式", "追溯"},
|
||||
"content": "后端需要组装系统 Prompt、检索片段、历史上下文和当前问题，通过 SSE 流式输出，并记录 AI 请求日志。",
|
||||
},
|
||||
]
|
||||
|
||||
@classmethod
|
||||
def retrieve(cls, question: str, scopes: list[KnowledgeScope]) -> list[RetrievedChunk]:
|
||||
if not scopes:
|
||||
return []
|
||||
|
||||
normalized_question = question.lower()
|
||||
matched_documents = []
|
||||
for document in cls._mock_documents:
|
||||
if any(keyword.lower() in normalized_question for keyword in document["keywords"]):
|
||||
matched_documents.append(document)
|
||||
|
||||
if not matched_documents:
|
||||
return []
|
||||
|
||||
primary_scope = scopes[0]
|
||||
return [
|
||||
RetrievedChunk(
|
||||
knowledge_id=primary_scope.id,
|
||||
knowledge_name=primary_scope.name,
|
||||
title=document["title"],
|
||||
content=document["content"],
|
||||
source_url=None,
|
||||
)
|
||||
for document in matched_documents[:3]
|
||||
]
|
||||
|
||||
|
||||
class PromptService:
|
||||
_default_prompt = (
|
||||
"你是企业飞书知识库 AI 助手。你只能基于提供的知识片段回答，不能编造。"
|
||||
"如果知识片段之间有冲突，需要说明不同来源的观点。"
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def build_prompt(cls, db: Session, question: str, chunks: list[RetrievedChunk]) -> str:
|
||||
prompt = cls._load_active_prompt(db)
|
||||
context = "\n\n".join(
|
||||
f"[知识片段 {index}] 来源：{chunk.knowledge_name} / {chunk.title}\n{chunk.content}"
|
||||
for index, chunk in enumerate(chunks, start=1)
|
||||
)
|
||||
if not context:
|
||||
context = "未检索到相关知识片段。"
|
||||
return f"{prompt}\n\n{context}\n\n用户问题：{question.strip()}"
|
||||
|
||||
@classmethod
|
||||
def _load_active_prompt(cls, db: Session) -> str:
|
||||
prompt = db.scalar(
|
||||
select(Prompt)
|
||||
.order_by(Prompt.updated_at.desc(), Prompt.id.desc())
|
||||
.limit(1)
|
||||
)
|
||||
return prompt.prompt_content if prompt else cls._default_prompt
|
||||
Reference in New Issue
Block a user