project/ex_module/ex_translated_profile_extractor.py

from typing import TypedDict

from ex_profile_extractor import ProfileExtractor, ProfileExtractorInputType
from ex_translator import Translator

from neollm import MyL3M2


class TranslatedProfileExtractorOutputType(TypedDict):
    name_ENG: str
    name_JPN: str
    domain_ENG: str
    domain_JPN: str
    birth_year: int


class TranslatedProfileExtractor(MyL3M2):
    def _link(self, inputs: ProfileExtractorInputType) -> TranslatedProfileExtractorOutputType:
        # Profile Extract
        profile_extractor = ProfileExtractor(parent=self, silent_list=["llm_settings", "inputs", "messages"])
        profile = profile_extractor(inputs)
        # Translator name
        translator_name = Translator(parent=self, silent_list=["llm_settings", "inputs", "messages"])
        translated_name = translator_name(inputs={"text": profile["name"]})["text_translated"]
        # Translate domain
        translator_domain = Translator(parent=self, silent_list=["llm_settings", "inputs", "messages"])
        translated_domain = translator_domain(inputs={"text": profile["domain"]})["text_translated"]

        outputs: TranslatedProfileExtractorOutputType = {
            "name_ENG": profile["name"],
            "name_JPN": profile["name"],
            "domain_ENG": profile["domain"],
            "domain_JPN": profile["domain"],
            "birth_year": profile["birth_year"],
        }

        if profile["lang"] == "ENG":
            outputs["name_JPN"] = translated_name
            outputs["domain_JPN"] = translated_domain
        else:
            outputs["name_ENG"] = translated_name
            outputs["domain_ENG"] = translated_domain

        return outputs

    # 型定義のために必要
    def __call__(self, inputs: ProfileExtractorInputType) -> TranslatedProfileExtractorOutputType:
        outputs: TranslatedProfileExtractorOutputType = super().__call__(inputs)
        return outputs