zhanghongyu_0626
70aa861983
- Added new modules for advisor compliance, KYC sessions, and script templates, enhancing the advisor agent's capabilities. - Implemented a comprehensive API structure under the `/api/advisor-agent` prefix, ensuring clear organization and access to new features. - Established database models and repositories for compliance rules and KYC sessions, facilitating robust data management. - Integrated exception handling and response models to improve error management and user feedback. - Updated settings to include new configurations for compliance and KYC features, ensuring flexibility and adaptability. This update significantly expands the advisor agent's functionality, providing essential tools for compliance and customer interaction while maintaining a structured API design.
131 lines
4.3 KiB
Python
131 lines
4.3 KiB
Python
"""Synchronize approved script templates to the Milvus vector collection."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from typing import Protocol
|
|
|
|
from app.config.settings import settings
|
|
from app.model.entities_advisor import ScriptTemplate
|
|
from app.repository.script_template_repository import ScriptTemplateRepository
|
|
from app.tool.embedding_tool import EmbeddingError, OllamaEmbeddingTool
|
|
from app.tool.milvus_template_tool import (
|
|
MilvusTemplateVectorStore,
|
|
MilvusToolError,
|
|
TemplateVectorHit,
|
|
TemplateVectorRecord,
|
|
)
|
|
|
|
|
|
class EmbeddingTool(Protocol):
|
|
def embed_text(self, text: str) -> list[float]:
|
|
pass
|
|
|
|
|
|
class TemplateVectorStore(Protocol):
|
|
def ensure_template_collection(self) -> None:
|
|
pass
|
|
|
|
def upsert_template(self, record: TemplateVectorRecord) -> str:
|
|
pass
|
|
|
|
def delete_template(self, template_id: int) -> None:
|
|
pass
|
|
|
|
def search_templates(
|
|
self,
|
|
*,
|
|
embedding: list[float],
|
|
scene: str | None = None,
|
|
top_k: int = 10,
|
|
) -> list[TemplateVectorHit]:
|
|
pass
|
|
|
|
|
|
class TemplateVectorError(RuntimeError):
|
|
pass
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class TemplateVectorSyncResult:
|
|
total: int
|
|
upserted: int
|
|
skipped: int
|
|
|
|
|
|
class ScriptTemplateVectorService:
|
|
def __init__(
|
|
self,
|
|
*,
|
|
repository: ScriptTemplateRepository | None = None,
|
|
embedding_tool: EmbeddingTool | None = None,
|
|
vector_store: TemplateVectorStore | None = None,
|
|
) -> None:
|
|
self.repository = repository or ScriptTemplateRepository()
|
|
self.embedding_tool = embedding_tool or OllamaEmbeddingTool()
|
|
self.vector_store = vector_store or MilvusTemplateVectorStore()
|
|
|
|
def sync_approved_templates(self, *, created_by: str | None = None) -> TemplateVectorSyncResult:
|
|
templates = self.repository.list_vector_candidates(created_by=created_by)
|
|
self.vector_store.ensure_template_collection()
|
|
upserted = 0
|
|
skipped = 0
|
|
for template in templates:
|
|
try:
|
|
self.upsert_template(template)
|
|
upserted += 1
|
|
except TemplateVectorError:
|
|
skipped += 1
|
|
return TemplateVectorSyncResult(total=len(templates), upserted=upserted, skipped=skipped)
|
|
|
|
def upsert_template(self, template: ScriptTemplate) -> str:
|
|
if not template.is_approved or not template.is_active:
|
|
self.delete_template(template.id)
|
|
return ""
|
|
|
|
try:
|
|
embedding = self.embedding_tool.embed_text(template.content)
|
|
except EmbeddingError as exc:
|
|
raise TemplateVectorError(str(exc)) from exc
|
|
if len(embedding) != settings.embed_dim:
|
|
raise TemplateVectorError(f"Embedding dimension must be {settings.embed_dim}, got {len(embedding)}")
|
|
vector_id = f"tpl_{template.id}_0"
|
|
try:
|
|
stored_id = self.vector_store.upsert_template(
|
|
TemplateVectorRecord(
|
|
vector_id=vector_id,
|
|
template_id=template.id,
|
|
embedding=embedding,
|
|
scene=template.scene,
|
|
title=template.title,
|
|
tags=";".join(template.tags or []),
|
|
is_approved=template.is_approved,
|
|
chunk_text=template.content,
|
|
chunk_no=0,
|
|
)
|
|
)
|
|
except MilvusToolError as exc:
|
|
raise TemplateVectorError(str(exc)) from exc
|
|
self.repository.update_embedding_id(template.id, stored_id, "template_vector_sync")
|
|
return stored_id
|
|
|
|
def delete_template(self, template_id: int) -> None:
|
|
try:
|
|
self.vector_store.delete_template(template_id)
|
|
except MilvusToolError as exc:
|
|
raise TemplateVectorError(str(exc)) from exc
|
|
self.repository.update_embedding_id(template_id, None, "template_vector_sync")
|
|
|
|
def search_templates(
|
|
self,
|
|
*,
|
|
query: str,
|
|
scene: str | None = None,
|
|
top_k: int = 10,
|
|
) -> list[TemplateVectorHit]:
|
|
try:
|
|
embedding = self.embedding_tool.embed_text(query)
|
|
return self.vector_store.search_templates(embedding=embedding, scene=scene, top_k=top_k)
|
|
except (EmbeddingError, MilvusToolError) as exc:
|
|
raise TemplateVectorError(str(exc)) from exc
|