231 lines
8.7 KiB
Python
231 lines
8.7 KiB
Python
import json
|
|
from typing import Any, Dict, Iterable, Mapping, Optional, Protocol
|
|
|
|
from fastapi import Depends, HTTPException
|
|
from sqlalchemy.orm import Session
|
|
|
|
from app.core.config import settings
|
|
from app.db import get_db_session
|
|
from shared_ai_platform_sdk import RequestContext as SharedAIRequestContext
|
|
from shared_ai_platform_sdk import SharedAIPlatformClient as SharedAIPlatformSdkClient
|
|
from shared_ai_platform_sdk import SharedAIPlatformError
|
|
|
|
|
|
class SharedAIPlatformClient(Protocol):
|
|
def stream_answer(
|
|
self,
|
|
current_user: Dict[str, Any],
|
|
conversation_id: Optional[str],
|
|
message: str,
|
|
review_context: Optional[Any] = None,
|
|
) -> Iterable[str]: ...
|
|
|
|
def list_knowledge_bases(self, current_user: Dict[str, Any], page: int, page_size: int) -> dict: ...
|
|
|
|
def list_conversations(self, current_user: Dict[str, Any], page: int, page_size: int) -> dict: ...
|
|
|
|
def get_conversation_detail(self, current_user: Dict[str, Any], conversation_id: str) -> dict: ...
|
|
|
|
def run_legal_review(self, current_user: Dict[str, Any], payload: Dict[str, Any]) -> dict: ...
|
|
|
|
def create_knowledge_base(self, current_user: Dict[str, Any], name: str, description: Optional[str]) -> dict: ...
|
|
|
|
def list_documents(self, current_user: Dict[str, Any], knowledge_base_id: str, page: int, page_size: int) -> dict: ...
|
|
|
|
def upload_document(
|
|
self,
|
|
current_user: Dict[str, Any],
|
|
knowledge_base_id: str,
|
|
file_name: str,
|
|
mime_type: Optional[str],
|
|
content: bytes,
|
|
) -> dict: ...
|
|
|
|
def retry_document(
|
|
self,
|
|
current_user: Dict[str, Any],
|
|
knowledge_base_id: str,
|
|
document_id: str,
|
|
file_name: Optional[str] = None,
|
|
mime_type: Optional[str] = None,
|
|
content: Optional[bytes] = None,
|
|
) -> dict: ...
|
|
|
|
|
|
|
|
class RemoteSharedAIPlatformClient:
|
|
def __init__(self, base_url: str, app_id: str, api_key: str, timeout_seconds: float):
|
|
self.base_url = base_url.rstrip("/")
|
|
self.app_id = app_id
|
|
self.api_key = api_key
|
|
self.timeout_seconds = timeout_seconds
|
|
self.sdk = SharedAIPlatformSdkClient(
|
|
base_url=self.base_url,
|
|
api_key=self.api_key,
|
|
app_id=self.app_id,
|
|
timeout_seconds=self.timeout_seconds,
|
|
)
|
|
|
|
def stream_answer(
|
|
self,
|
|
current_user: Dict[str, Any],
|
|
conversation_id: Optional[str],
|
|
message: str,
|
|
review_context: Optional[Any] = None,
|
|
) -> Iterable[str]:
|
|
review_context_payload = self._normalize_review_context(review_context)
|
|
|
|
def _event_stream():
|
|
try:
|
|
for event in self.sdk.stream_conversation(
|
|
self._context(current_user),
|
|
message=message,
|
|
chat_context=review_context_payload or {},
|
|
conversation_id=conversation_id,
|
|
):
|
|
if event.event_type in {"delta", "text"}:
|
|
yield self._format_event({"type": "text", "content": event.payload.get("content", "")})
|
|
elif event.event_type == "sources":
|
|
yield self._format_event({"type": "sources", "sources": event.payload.get("sources", [])})
|
|
elif event.event_type == "end":
|
|
yield self._format_event(
|
|
{
|
|
"type": "end",
|
|
"conversation_id": event.payload.get("conversation_id"),
|
|
}
|
|
)
|
|
elif event.event_type == "error":
|
|
yield self._format_event({"type": "error", "message": event.payload.get("message", "shared ai platform request failed")})
|
|
elif event.event_type in {"start", "usage"}:
|
|
continue
|
|
except SharedAIPlatformError as exc:
|
|
raise HTTPException(status_code=exc.status_code, detail=str(exc)) from exc
|
|
|
|
return _event_stream()
|
|
|
|
def list_knowledge_bases(self, current_user: Dict[str, Any], page: int, page_size: int) -> dict:
|
|
payload = self.sdk.list_knowledge_bases(self._context(current_user))
|
|
return self._apply_remote_knowledge_defaults(payload)
|
|
|
|
def list_conversations(self, current_user: Dict[str, Any], page: int, page_size: int) -> dict:
|
|
return self.sdk.list_conversations(self._context(current_user), page=page, page_size=page_size)
|
|
|
|
def get_conversation_detail(self, current_user: Dict[str, Any], conversation_id: str) -> dict:
|
|
return self.sdk.get_conversation_detail(self._context(current_user), conversation_id)
|
|
|
|
def run_legal_review(self, current_user: Dict[str, Any], payload: Dict[str, Any]) -> dict:
|
|
return self.sdk.run_legal_review(
|
|
self._context(current_user),
|
|
title=payload.get("title") or "文本输入合同",
|
|
text=payload.get("text") or "",
|
|
contract_type=payload.get("contract_type"),
|
|
review_as=payload.get("review_as") or "party_a",
|
|
provider=payload.get("provider"),
|
|
)
|
|
|
|
def create_knowledge_base(self, current_user: Dict[str, Any], name: str, description: Optional[str]) -> dict:
|
|
payload = self.sdk.create_knowledge_base(
|
|
self._context(current_user),
|
|
name=name,
|
|
description=description,
|
|
retrieval_mode="hybrid",
|
|
)
|
|
if "retrieval_mode" not in payload:
|
|
payload["retrieval_mode"] = "hybrid"
|
|
return payload
|
|
|
|
def list_documents(self, current_user: Dict[str, Any], knowledge_base_id: str, page: int, page_size: int) -> dict:
|
|
return self.sdk.list_documents(self._context(current_user), knowledge_base_id, page=page, page_size=page_size)
|
|
|
|
def upload_document(
|
|
self,
|
|
current_user: Dict[str, Any],
|
|
knowledge_base_id: str,
|
|
file_name: str,
|
|
mime_type: Optional[str],
|
|
content: bytes,
|
|
) -> dict:
|
|
return self.sdk.upload_document(
|
|
self._context(current_user),
|
|
knowledge_base_id,
|
|
file_name=file_name,
|
|
mime_type=mime_type,
|
|
content=content,
|
|
)
|
|
|
|
def retry_document(
|
|
self,
|
|
current_user: Dict[str, Any],
|
|
knowledge_base_id: str,
|
|
document_id: str,
|
|
file_name: Optional[str] = None,
|
|
mime_type: Optional[str] = None,
|
|
content: Optional[bytes] = None,
|
|
) -> dict:
|
|
return self.sdk.retry_document(
|
|
self._context(current_user),
|
|
knowledge_base_id,
|
|
document_id,
|
|
file_name=file_name,
|
|
mime_type=mime_type,
|
|
content=content,
|
|
)
|
|
|
|
def _context(self, current_user: Dict[str, Any]) -> SharedAIRequestContext:
|
|
return SharedAIRequestContext(
|
|
tenant_id=str(current_user.get("tenant_id") or ""),
|
|
user_id=str(current_user.get("id") or ""),
|
|
user_role=str(current_user.get("role") or "admin"),
|
|
)
|
|
|
|
def _format_event(self, payload: Dict[str, object]) -> str:
|
|
return f"data: {json.dumps(payload, ensure_ascii=False)}\n\n"
|
|
|
|
def _apply_remote_knowledge_defaults(self, payload: dict) -> dict:
|
|
items = payload.get("items")
|
|
if not isinstance(items, list):
|
|
return payload
|
|
|
|
for item in items:
|
|
if isinstance(item, dict) and not item.get("retrieval_mode"):
|
|
item["retrieval_mode"] = "hybrid"
|
|
return payload
|
|
|
|
def _normalize_review_context(self, review_context: Any) -> Optional[Dict[str, Any]]:
|
|
if review_context is None:
|
|
return None
|
|
if isinstance(review_context, Mapping):
|
|
return dict(review_context)
|
|
if hasattr(review_context, "model_dump"):
|
|
return review_context.model_dump(mode="json", by_alias=True, exclude_none=True)
|
|
return review_context
|
|
|
|
|
|
|
|
def get_shared_ai_platform_client(db: Session = Depends(get_db_session)) -> SharedAIPlatformClient:
|
|
return build_shared_ai_platform_client(db)
|
|
|
|
|
|
|
|
def build_shared_ai_platform_client(db: Session) -> SharedAIPlatformClient:
|
|
base_url = settings.shared_ai_platform_base_url.strip()
|
|
if not base_url:
|
|
raise HTTPException(
|
|
status_code=503,
|
|
detail="shared-ai-platform base URL is required for AI features but is not configured",
|
|
)
|
|
|
|
api_key = settings.shared_ai_platform_api_key.strip()
|
|
if not api_key:
|
|
raise HTTPException(
|
|
status_code=503,
|
|
detail="shared-ai-platform API key is required for AI features but is not configured",
|
|
)
|
|
|
|
return RemoteSharedAIPlatformClient(
|
|
base_url=base_url,
|
|
app_id=settings.shared_ai_platform_app_id,
|
|
api_key=api_key,
|
|
timeout_seconds=settings.shared_ai_platform_timeout_seconds,
|
|
)
|