dieya_risk_prediction_backend/app/repositories/integration_conversations.py

318 lines
13 KiB
Python

"""Durable report-to-chat links and idempotent external chat turns."""
import hashlib
import uuid
from datetime import UTC, datetime, timedelta
from sqlalchemy import select, text, update
from sqlalchemy.dialects.postgresql import insert
from sqlalchemy.ext.asyncio import AsyncSession
from app.core.exceptions import AppException
from app.models.entities import (
Conversation,
ConversationMessage,
IntegrationAssessment,
IntegrationConversation,
IntegrationConversationRequest,
OperationAudit,
)
class IntegrationConversationRepository:
def __init__(self, session: AsyncSession) -> None:
self.session = session
async def lock_conversation(self, client_id: str, conversation_id: int) -> None:
lock_key = int.from_bytes(
hashlib.sha256(f"chat-conversation:{client_id}:{conversation_id}".encode()).digest()[:8],
signed=True,
)
await self.session.execute(text("SELECT pg_advisory_xact_lock(:key)"), {"key": lock_key})
async def get_report(self, client_id: str, request_id: str) -> IntegrationAssessment | None:
return await self.session.scalar(
select(IntegrationAssessment).where(
IntegrationAssessment.client_id == client_id,
IntegrationAssessment.request_id == request_id,
)
)
async def get_by_report(self, client_id: str, request_id: str) -> IntegrationConversation | None:
return await self.session.scalar(
select(IntegrationConversation).where(
IntegrationConversation.client_id == client_id,
IntegrationConversation.report_request_id == request_id,
)
)
async def get_conversation(self, conversation_id: int, user_id: int) -> Conversation | None:
return await self.session.scalar(
select(Conversation).where(
Conversation.id == conversation_id,
Conversation.user_id == user_id,
)
)
async def create_or_get(
self,
*,
client_id: str,
report: IntegrationAssessment,
title: str,
actor_emp_no: str,
) -> tuple[IntegrationConversation, Conversation, bool]:
lock_key = int.from_bytes(
hashlib.sha256(f"report-chat:{client_id}:{report.request_id}".encode()).digest()[:8],
signed=True,
)
await self.session.execute(text("SELECT pg_advisory_xact_lock(:key)"), {"key": lock_key})
link = await self.get_by_report(client_id, report.request_id)
if link is not None:
conversation = await self.get_conversation(link.conversation_id, link.user_id)
if conversation is None:
raise AppException("报告会话记录无效", code=500, status_code=500)
return link, conversation, False
now = datetime.now(UTC)
conversation = Conversation(user_id=report.user_id, title=title[:120])
self.session.add(conversation)
await self.session.flush()
link = IntegrationConversation(
client_id=client_id,
report_request_id=report.request_id,
conversation_id=conversation.id,
user_id=report.user_id,
created_at=now,
)
self.session.add(link)
self.session.add(
OperationAudit(
actor_user_id=report.user_id,
actor_emp_no=actor_emp_no,
action="integration.conversation.create",
resource_type="conversation",
resource_id=str(conversation.id),
audit_metadata={"clientId": client_id, "reportRequestId": report.request_id},
)
)
await self.session.flush()
return link, conversation, True
async def claim_message(
self,
*,
client_id: str,
conversation_id: int,
request_id: str,
user_id: int,
content: str,
) -> tuple[IntegrationConversationRequest, bool]:
await self.lock_conversation(client_id, conversation_id)
digest = hashlib.sha256(content.encode("utf-8")).hexdigest()
now = datetime.now(UTC)
token = uuid.uuid4().hex
lease_until = now + timedelta(minutes=5)
existing = await self.get_message_request(client_id, conversation_id, request_id)
if existing is not None:
if existing.content_hash != digest or existing.user_id != user_id:
raise AppException("messageRequestId 已用于其他用户或不同内容", code=409, status_code=409)
expires = existing.lease_until
if expires.tzinfo is None:
expires = expires.replace(tzinfo=UTC)
if existing.status == "completed" or (existing.status == "processing" and expires > now):
await self.session.commit()
return existing, False
in_flight = await self.session.scalar(
select(IntegrationConversationRequest.request_id)
.where(
IntegrationConversationRequest.client_id == client_id,
IntegrationConversationRequest.conversation_id == conversation_id,
IntegrationConversationRequest.request_id != request_id,
IntegrationConversationRequest.status.in_(["processing", "failed"]),
)
.limit(1)
)
if in_flight is not None:
raise AppException("会话存在未完成追问,请先恢复原 messageRequestId", code=409, status_code=409)
inserted = await self.session.scalar(
insert(IntegrationConversationRequest)
.values(
client_id=client_id,
conversation_id=conversation_id,
request_id=request_id,
user_id=user_id,
content_hash=digest,
status="processing",
owner_token=token,
lease_until=lease_until,
created_at=now,
updated_at=now,
)
.on_conflict_do_nothing(index_elements=["client_id", "conversation_id", "request_id"])
.returning(IntegrationConversationRequest.id)
)
record = await self.get_message_request(client_id, conversation_id, request_id)
if record is None:
raise AppException("追问请求记录创建失败", code=503, status_code=503)
if record.content_hash != digest or record.user_id != user_id:
raise AppException("messageRequestId 已用于其他用户或不同内容", code=409, status_code=409)
if inserted is not None:
await self._attach_user_message(record, token, user_id, content, now)
return record, True
# Reclaim only an interrupted/failed request. The old worker's token can
# no longer persist a reply after this transaction commits.
reclaimed = await self.session.scalar(
update(IntegrationConversationRequest)
.where(
IntegrationConversationRequest.id == record.id,
IntegrationConversationRequest.status != "completed",
(IntegrationConversationRequest.status == "failed")
| (IntegrationConversationRequest.lease_until <= now),
)
.values(
status="processing",
owner_token=token,
lease_until=lease_until,
error_message=None,
updated_at=now,
)
.execution_options(synchronize_session=False)
.returning(IntegrationConversationRequest.id)
)
await self.session.commit()
record = await self.get_message_request(client_id, conversation_id, request_id)
if reclaimed is not None and record.user_message_id is None:
await self._attach_user_message(record, token, user_id, content, datetime.now(UTC))
record = await self.get_message_request(client_id, conversation_id, request_id)
return record, reclaimed is not None
async def _attach_user_message(
self,
record: IntegrationConversationRequest,
token: str,
user_id: int,
content: str,
now: datetime,
) -> None:
conversation = await self.get_conversation(record.conversation_id, user_id)
if conversation is None:
raise AppException("报告会话不存在或无权访问", code=404, status_code=404)
message = ConversationMessage(
conversation_id=conversation.id,
role="user",
message_type="text",
content=content,
structured_data={"integrationRequestId": record.request_id},
created_at=now,
)
self.session.add(message)
conversation.updated_at = now
await self.session.flush()
await self.session.execute(
update(IntegrationConversationRequest)
.where(
IntegrationConversationRequest.id == record.id,
IntegrationConversationRequest.owner_token == token,
IntegrationConversationRequest.status == "processing",
)
.values(user_message_id=message.id, updated_at=now)
)
await self.session.commit()
record.user_message_id = message.id
async def get_message_request(
self,
client_id: str,
conversation_id: int,
request_id: str,
) -> IntegrationConversationRequest | None:
return await self.session.scalar(
select(IntegrationConversationRequest)
.where(
IntegrationConversationRequest.client_id == client_id,
IntegrationConversationRequest.conversation_id == conversation_id,
IntegrationConversationRequest.request_id == request_id,
)
.execution_options(populate_existing=True)
)
async def get_conversation_message(
self,
conversation_id: int,
message_id: int,
) -> ConversationMessage | None:
return await self.session.scalar(
select(ConversationMessage).where(
ConversationMessage.conversation_id == conversation_id,
ConversationMessage.id == message_id,
)
)
async def persist_reply(
self,
*,
client_id: str,
conversation_id: int,
request_id: str,
token: str,
content: str,
semantic: dict,
actor_emp_no: str,
) -> ConversationMessage:
await self.lock_conversation(client_id, conversation_id)
record = await self.get_message_request(client_id, conversation_id, request_id)
now = datetime.now(UTC)
expires = record.lease_until if record else now
if expires.tzinfo is None:
expires = expires.replace(tzinfo=UTC)
if record is None or record.status != "processing" or record.owner_token != token or expires <= now:
await self.session.rollback()
raise AppException("追问处理权已过期,请查询原请求状态后重试", code=409, status_code=409)
message = ConversationMessage(
conversation_id=conversation_id,
role="assistant",
message_type="text",
content=content,
structured_data={"semantic": semantic, "integrationRequestId": request_id},
created_at=now,
)
self.session.add(message)
await self.session.flush()
record.status = "completed"
record.assistant_message_id = message.id
record.updated_at = now
record.error_message = None
await self.session.execute(
update(Conversation).where(Conversation.id == conversation_id).values(updated_at=now)
)
self.session.add(
OperationAudit(
actor_user_id=record.user_id,
actor_emp_no=actor_emp_no,
action="integration.conversation.reply",
resource_type="conversation",
resource_id=str(conversation_id),
audit_metadata={"clientId": client_id, "messageRequestId": request_id},
)
)
await self.session.commit()
return message
async def fail_message(self, record_id: int, token: str, error: str) -> None:
await self.session.execute(
update(IntegrationConversationRequest)
.where(
IntegrationConversationRequest.id == record_id,
IntegrationConversationRequest.owner_token == token,
IntegrationConversationRequest.status == "processing",
)
.values(
status="failed",
error_message=error[:200],
lease_until=datetime.now(UTC),
updated_at=datetime.now(UTC),
)
)
await self.session.commit()