""" EasyRead Controller. Orchestrates the EasyRead pipeline (simplify -> validate -> revise -> icons) over the services in ``src/services`` and the prompts in ``src/prompts``. Provider selection ------------------ * LLM provider (``llm_provider``): which model drives simplification, validation and revision. Defaults to ``"llama"`` (Meta Llama 3 70B on AWS Bedrock). Alternatives: ``"gemini"``, ``"bedrock"``. * Translation provider (``translation_provider``): how sentences are translated. Defaults to ``"nllb"`` (facebook/nllb-200). Alternatives: ``"google"`` (Google Translate) and ``"gemini"``. Heavy or credential-bound dependencies (the icon generator, translators) are loaded lazily so constructing the Controller stays cheap and side-effect free. Usage: from controller import Controller controller = Controller() # llama + nllb (defaults) controller = Controller(llm_provider="gemini", translation_provider="google") result = controller.simplify_text(markdown, target_language="sw") """ import json import re import sys from logging import getLogger from pathlib import Path from uuid import uuid4 import yaml # Make ``src`` importable whether the controller is run from the repo root or # imported as a module. BASE_DIR = Path(__file__).resolve().parent if str(BASE_DIR) not in sys.path: sys.path.insert(0, str(BASE_DIR)) logger = getLogger(__name__) PROMPTS_DIR = BASE_DIR / "src" / "prompts" ICON_OUTPUT_PATH = BASE_DIR / "src" / "temp" # Supported providers (used for validation / friendly errors). LLM_PROVIDERS = ("hf", "llama", "gemini", "bedrock") TRANSLATION_PROVIDERS = ("nllb", "google", "gemini") def _strip_code_fences(text: str) -> str: """Remove markdown code fences (```json ... ```) that LLMs sometimes add.""" match = re.search(r'```(?:json)?\s*([\s\S]*?)```', text) return match.group(1).strip() if match else text.strip() class Controller: """Drive the EasyRead pipeline with pluggable LLM and translation backends.""" def __init__( self, llm_provider: str = "hf", translation_provider: str = "nllb", ): """ Args: llm_provider: One of ``LLM_PROVIDERS``. Default ``"hf"`` (Llama via Hugging Face Inference Providers). translation_provider: One of ``TRANSLATION_PROVIDERS``. Default ``"nllb"``. """ self.llm_provider = llm_provider.lower() self.translation_provider = translation_provider.lower() if self.llm_provider not in LLM_PROVIDERS: raise ValueError( f"Unknown llm_provider '{llm_provider}'. Choose from: {', '.join(LLM_PROVIDERS)}" ) if self.translation_provider not in TRANSLATION_PROVIDERS: raise ValueError( f"Unknown translation_provider '{translation_provider}'. " f"Choose from: {', '.join(TRANSLATION_PROVIDERS)}" ) # Eagerly build the selected LLM (the pipeline's core dependency). self.llm = self._make_llm(self.llm_provider) # Prompts loaded from src/prompts/*.yaml. self.prompts = self._load_prompts() self.icon_output_path = ICON_OUTPUT_PATH # Lazily-initialised collaborators. self._translator = None self._gemini = None self._icon_generator = None self._global_symbols = None # ------------------------------------------------------------------ # # Provider factories # # ------------------------------------------------------------------ # @staticmethod def _make_llm(provider: str): """Instantiate the chosen LLM driver. Each exposes generate_text(prompt).""" if provider == "hf": from src.services.hf_inference import HFInferenceDriver return HFInferenceDriver() if provider == "llama": from src.services.llama_3_70b import LlamaDriver return LlamaDriver() if provider == "gemini": from src.services.gemini import GeminiDriver return GeminiDriver() if provider == "bedrock": from src.services.bedrock import BedrockDriver return BedrockDriver() raise ValueError(f"Unknown llm_provider '{provider}'") def _get_translator(self): """Lazily build the chosen translation backend (google or nllb).""" if self._translator is not None: return self._translator if self.translation_provider == "google": from src.services.google_translate import GoogleTranslateDriver self._translator = GoogleTranslateDriver() elif self.translation_provider == "nllb": try: from src.services.nllb_200 import NLLBDriver except ImportError as exc: raise RuntimeError( "NLLB translation driver (src/services/nllb_200.py) is not " "implemented yet. Provide the NLLBDriver script, or construct " "the Controller with translation_provider='google' or 'gemini'." ) from exc self._translator = NLLBDriver() else: # pragma: no cover - gemini handled in translate_text raise RuntimeError(f"No standalone translator for '{self.translation_provider}'") return self._translator def _get_gemini(self): """Lazily build a GeminiDriver (used for gemini-based translation).""" if self._gemini is None: from src.services.gemini import GeminiDriver # Reuse the LLM instance if it already is Gemini. self._gemini = self.llm if self.llm_provider == "gemini" else GeminiDriver() return self._gemini @property def icon_generator(self): """Lazily load the ETH LoRA icon generator (pulls torch/diffusers).""" if self._icon_generator is None: from src.services.eth_zurich import IconGenerator self._icon_generator = IconGenerator() return self._icon_generator @property def global_symbols(self): """Lazily build the Global Symbols API client.""" if self._global_symbols is None: from src.services.global_symbols import GlobalSymbolsService self._global_symbols = GlobalSymbolsService() return self._global_symbols # ------------------------------------------------------------------ # # Prompt loading # # ------------------------------------------------------------------ # @staticmethod def _load_prompts() -> dict: """Load all prompt YAMLs from src/prompts keyed by filename stem.""" prompts = {} for path in PROMPTS_DIR.glob("*.yaml"): with open(path) as f: prompts[path.stem] = yaml.safe_load(f) or {} logger.info("Loaded prompts: %s", ", ".join(sorted(prompts))) return prompts # ------------------------------------------------------------------ # # Translation # # ------------------------------------------------------------------ # def translate_text(self, text: str, target_language: str) -> str: """ Translate text using the configured translation provider. Args: text: Text to translate. target_language: Target language (ISO code for nllb/google). Returns: The translated text (or the original if ``target_language`` is empty). """ if not target_language: return text if self.translation_provider == "gemini": gemini = self._get_gemini() prompt = ( f"Translate the following text to {target_language}. " "Output only the translation, with no extra commentary or quotes.\n\n" f"Text:\n{text}" ) return gemini.generate_text(prompt).strip() # google / nllb expose translate(text, target_language). return self._get_translator().translate(text, target_language) # ------------------------------------------------------------------ # # Core pipeline # # ------------------------------------------------------------------ # def simplify_text(self, text: str, target_language: str = None) -> dict: """Convert markdown into Easy Read sentences, optionally translated.""" template = self.prompts['simplify_text']['system_message'] prompt = f"{template}\n\n Bellow is the Input Text to simplify:\n\n{text}\n\n" response = self.llm.generate_text(prompt) try: response_data = json.loads(_strip_code_fences(response)) logger.info("Successfully parsed simplify response JSON") if target_language: for s in response_data.get('simplified_sentences', []): s['translated_sentence'] = self.translate_text(s['sentence'], target_language) except json.JSONDecodeError: logger.error(f"Failed to parse response as JSON. Raw response: {response}") response_data = {"error": "Failed to parse response as JSON.", "raw_response": response} return response_data def validate_text(self, original_sentence: str, simplified_sentences: list[dict]) -> dict: """Validate Easy Read sentences against the original markdown.""" template = self.prompts['validate_text']['system_message'] prompt = template + "\n" + self.prompts['validate_text']["user_message_template"].format( original_markdown=original_sentence, simplified_sentences=json.dumps(simplified_sentences), ) response = self.llm.generate_text(prompt) try: response_data = json.loads(_strip_code_fences(response)) logger.info("Successfully parsed validate response JSON") except json.JSONDecodeError: response_data = {"error": "Failed to parse response as JSON.", "raw_response": response} logger.error(f"Failed to parse response as JSON. Raw response: {response}") return response_data def revise_text( self, original_text: str, easy_read_sentences: list, feedback: str, target_language: str = None, ) -> dict: """Revise Easy Read sentences using validation feedback, optionally translated.""" template = self.prompts['revise_text']['system_message'] prompt = template + "\n" + self.prompts['revise_text']["user_message_template"].format( original_markdown=original_text, simplified_sentences=json.dumps(easy_read_sentences), validation_feedback=feedback, ) response = self.llm.generate_text(prompt) try: response_data = json.loads(_strip_code_fences(response)) logger.info("Successfully parsed revise response JSON") if target_language: for s in response_data.get('revised_sentences', []): s['translated_sentence'] = self.translate_text(s['sentence'], target_language) except json.JSONDecodeError: response_data = {"error": "Failed to parse response as JSON.", "raw_response": response} logger.error(f"Failed to parse response as JSON. Raw response: {response}") return response_data # ------------------------------------------------------------------ # # Icons / symbols # # ------------------------------------------------------------------ # def generate_icons( self, sentences: list[dict], symbolset: str = "arasaac", use_global_symbols: bool = True, ) -> dict: """ Generate icons for sentences using Global Symbols API, falling back to the local LoRA generator when no symbol is found. """ request_id = str(uuid4()) request_dir = Path(self.icon_output_path) / request_id request_dir.mkdir(parents=True, exist_ok=True) for sentence in sentences: prompt = sentence['image_prompt'] safe_prompt = "_".join(prompt.split()) image_path = request_dir / f"{safe_prompt}.png" image_found = False if use_global_symbols: logger.info(f"Querying Global Symbols API for prompt: '{prompt}'") try: downloaded_path = self.global_symbols.search_and_download( query=prompt, output_path=image_path, symbolset=symbolset, ) if downloaded_path: logger.info(f"Retrieved symbol from Global Symbols for '{prompt}'") image_found = True else: logger.warning(f"No symbol found in Global Symbols for '{prompt}'") except Exception as e: logger.error(f"Error querying Global Symbols for '{prompt}': {e}") if not image_found: logger.info(f"Generating icon locally for prompt: '{prompt}'") try: image = self.icon_generator.generate(prompt) image.save(image_path) logger.info(f"Saved locally generated icon for '{prompt}'") except Exception as e: logger.error(f"Error generating icon locally for '{prompt}': {e}") continue sentence['image_path'] = "/".join(["icons", request_id, f"{safe_prompt}.png"]) return {"request_id": request_id, "icons": sentences} def search_symbols(self, sentences: list[dict], symbolset: str = "arasaac") -> dict: """Search Global Symbols for each sentence without AI fallback.""" request_id = str(uuid4()) request_dir = Path(self.icon_output_path) / request_id request_dir.mkdir(parents=True, exist_ok=True) results = [] for sentence in sentences: prompt = sentence['image_prompt'] safe_prompt = "_".join(prompt.split()) image_path = request_dir / f"{safe_prompt}.png" symbol_found = False symbol_image_path = None try: downloaded_path = self.global_symbols.search_and_download( query=prompt, output_path=image_path, symbolset=symbolset, ) if downloaded_path: symbol_found = True symbol_image_path = "/".join(["icons", request_id, f"{safe_prompt}.png"]) logger.info(f"Found symbol for '{prompt}'") else: logger.info(f"No symbol found for '{prompt}'") except Exception as e: logger.error(f"Error searching symbols for '{prompt}': {e}") results.append({ "sentence": sentence['sentence'], "image_prompt": sentence['image_prompt'], "highlighted": sentence['highlighted'], "symbol_found": symbol_found, "symbol_image_path": symbol_image_path, }) return {"request_id": request_id, "results": results} def generate_ai_icons(self, request_id: str, sentences: list[dict]) -> dict: """ Generate AI icons for sentences, passing through symbol search results where available and generating the rest with the ETH LoRA model. """ request_dir = Path(self.icon_output_path) / request_id request_dir.mkdir(parents=True, exist_ok=True) for sentence in sentences: symbol_image_path = sentence.get('symbol_image_path') if symbol_image_path: sentence['image_path'] = symbol_image_path logger.info(f"Using existing symbol for '{sentence['sentence']}'") else: ai_prompt = sentence['ai_prompt'] safe_prompt = "_".join(ai_prompt.split())[:80] image_path = request_dir / f"ai_{safe_prompt}.png" try: image = self.icon_generator.generate(ai_prompt) image.save(image_path) logger.info(f"Saved AI-generated icon for '{ai_prompt}'") except Exception as e: logger.error(f"Error generating AI icon for '{ai_prompt}': {e}") continue sentence['image_path'] = "/".join(["icons", request_id, f"ai_{safe_prompt}.png"]) sentence['image_prompt'] = sentence.get('ai_prompt', sentence.get('sentence', '')) return {"request_id": request_id, "icons": sentences} if __name__ == "__main__": import argparse parser = argparse.ArgumentParser(description="EasyRead controller — simplify markdown into Easy Read") parser.add_argument("text", type=str, help="Markdown/text to simplify") parser.add_argument("--llm", type=str, default="hf", choices=LLM_PROVIDERS, help="LLM provider") parser.add_argument("--translation", type=str, default="nllb", choices=TRANSLATION_PROVIDERS, help="Translation provider") parser.add_argument("--target-language", type=str, default=None, help="Optional target language code (e.g. sw)") args = parser.parse_args() controller = Controller(llm_provider=args.llm, translation_provider=args.translation) output = controller.simplify_text(args.text, target_language=args.target_language) print(json.dumps(output, ensure_ascii=False, indent=2))