LongjoeAgent/core/document_processor/templates/template_registry.py

from typing import Dict, Type
from core.document_processor.templates.base_template import DocumentTemplate
from core.document_processor.templates.invoice_template import InvoiceTemplate
from core.document_processor.templates.saletask_template import SaleTaskTemplate
from core.document_processor.templates.cusamt_template import CusAmtTemplate
from core.document_processor.templates.bank_template import BankTemplate


class TemplateRegistry:
    """模板注册管理器"""

    _templates: Dict[str, Type[DocumentTemplate]] = {}

    @classmethod
    def register(cls, template_class: Type[DocumentTemplate]):
        """注册模板"""
        instance = template_class()
        cls._templates[instance.template_name] = template_class
        return template_class

    @classmethod
    def get_template(cls, template_name: str) -> DocumentTemplate:
        """获取模板实例"""
        if template_name not in cls._templates:
            raise ValueError(f"模板不存在: {template_name}")
        return cls._templates[template_name]()

    @classmethod
    def list_templates(cls) -> Dict[str, str]:
        """列出所有可用模板"""
        return {name: cls._templates[name]().description for name in cls._templates}


# 注册内置模板
TemplateRegistry.register(InvoiceTemplate)
TemplateRegistry.register(SaleTaskTemplate)
TemplateRegistry.register(CusAmtTemplate)
TemplateRegistry.register(BankTemplate)