Update translate.py

aphra/translate.py

@@ -6,6 +6,28 @@
 from .prompts import get_prompt
 from .parsers import parse_analysis, parse_translation
 
+class TranslationContext:
+    """
+    Context for translation containing parameters and settings.
+
+    This class encapsulates the parameters and settings needed for performing a translation,
+    including the model client, source and target languages, and logging preferences.
+    """
+
+    def __init__(self, model_client, source_language, target_language, log_calls):
+        """
+        Initializes the TranslationContext with the provided parameters.
+
+        :param model_client: An instance of LLMModelClient.
+        :param source_language: The source language of the text.
+        :param target_language: The target language of the text.
+        :param log_calls: Boolean indicating whether to log the call details.
+        """
+        self.model_client = model_client
+        self.source_language = source_language
+        self.target_language = target_language
+        self.log_calls = log_calls
+
 def load_model_client(config_file):
     """
     Loads the LLMModelClient with the provided configuration file.
@@ -15,46 +37,41 @@ def load_model_client(config_file):
     """
     return LLMModelClient(config_file)
 
-def execute_model_call(model_client, system_file, user_file, model_name, log_calls, **kwargs):
+def execute_model_call(context, system_file, user_file, model_name, **kwargs):
     """
     Executes a model call using the provided system and user prompts.
 
-    :param model_client: An instance of LLMModelClient.
+    :param context: An instance of TranslationContext containing translation parameters.
     :param system_file: Path to the file containing the system prompt.
     :param user_file: Path to the file containing the user prompt.
     :param model_name: The name of the model to use.
-    :param log_calls: Boolean indicating whether to log the call details.
     :param kwargs: Optional keyword arguments to format the prompt templates.
     :return: The model's response content.
     """
     system_prompt = get_prompt(system_file, **kwargs)
     user_prompt = get_prompt(user_file, **kwargs)
-    return model_client.call_model(system_prompt, user_prompt, model_name, log_call=log_calls)
+    return context.model_client.call_model(system_prompt, user_prompt, model_name, log_call=context.log_calls)
 
-def generate_glossary(model_client, parsed_items, source_language, target_language, model_searcher, log_calls):
+def generate_glossary(context, parsed_items, model_searcher):
     """
     Generates a glossary of terms based on the parsed analysis items.
 
-    :param model_client: An instance of LLMModelClient.
+    :param context: An instance of TranslationContext containing translation parameters.
     :param parsed_items: A list of dictionaries containing 'name' and 'keywords' for each item.
-    :param source_language: The source language of the text.
-    :param target_language: The target language of the text.
     :param model_searcher: The name of the model to use for searching term explanations.
-    :param log_calls: Boolean indicating whether to log the call details.
     :return: A formatted string containing the glossary entries.
     """
     glossary = []
     for item in parsed_items:
         term_explanation = execute_model_call(
-            model_client,
+            context,
             'step2_system.txt',
             'step2_user.txt',
             model_searcher,
-            log_calls,
             term=item['name'],
             keywords=", ".join(item['keywords']),
-            source_language=source_language,
-            target_language=target_language
+            source_language=context.source_language,
+            target_language=context.target_language
         )
         glossary_entry = (
             f"### {item['name']}\n\n**Keywords:** {', '.join(item['keywords'])}\n\n"
@@ -76,40 +93,38 @@ def translate(source_language, target_language, text, config_file="config.toml",
     """
     model_client = load_model_client(config_file)
     models = model_client.llms
+    context = TranslationContext(model_client, source_language, target_language, log_calls)
 
     analysis_content = execute_model_call(
-        model_client,
+        context,
         'step1_system.txt',
         'step1_user.txt',
         models['writer'],
-        log_calls,
         post_content=text,
         source_language=source_language,
         target_language=target_language
     )
 
     parsed_items = parse_analysis(analysis_content)
     glossary_content = generate_glossary(
-        model_client, parsed_items, source_language, target_language, models['searcher'], log_calls
+        context, parsed_items, models['searcher']
     )
 
     translated_content = execute_model_call(
-        model_client,
+        context,
         'step3_system.txt',
         'step3_user.txt',
         models['writer'],
-        log_calls,
         text=text,
         source_language=source_language,
         target_language=target_language
     )
 
     critique = execute_model_call(
-        model_client,
+        context,
         'step4_system.txt',
         'step4_user.txt',
         models['critiquer'],
-        log_calls,
         text=text,
         translation=translated_content,
         glossary=glossary_content,
@@ -118,11 +133,10 @@ def translate(source_language, target_language, text, config_file="config.toml",
     )
 
     final_translation_content = execute_model_call(
-        model_client,
+        context,
         'step5_system.txt',
         'step5_user.txt',
         models['writer'],
-        log_calls,
         text=text,
         translation=translated_content,
         glossary=glossary_content,