diff --git a/main.py b/main.py index 95d17f5..d10c0b4 100644 --- a/main.py +++ b/main.py @@ -12,10 +12,54 @@ import zipfile import tempfile from openpyxl.utils import get_column_letter from openpyxl.styles import PatternFill, Border, Side, Alignment, Protection, Font - +from pydantic_ai import Agent +from pydantic_ai.models.openai import OpenAIModel +from pydantic_ai.providers.openai import OpenAIProvider +from pydantic import BaseModel, Field +from typing import Optional # Configure logging -logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') +logging.basicConfig(level=None, format='%(asctime)s - %(levelname)s - %(message)s') + + +class TranslationOutput(BaseModel): + translated_text: str +async def translate_with_ollama(text: str, target_language: str, model_name: str = "llama3.1:latest") -> str: + """Translate text using Ollama via PydanticAI + + Args: + text: Text to translate + target_language: Target language code (e.g. 'fr', 'en') + model_name: Ollama model name (default: llama3:latest) + + Returns: + Translated text + """ + # Skip empty text + if not text or text.strip() == "": + return text + + try: + # Create agent with Ollama + ollama_model = OpenAIModel( + model_name=model_name, provider=OpenAIProvider(base_url='http://localhost:11434/v1') + ) + agent = Agent(ollama_model, result_type=TranslationOutput,retries=2) + # Create the input + + + # Execute translation + result = await agent.run( + f"Translate this text to this language: {text} to {target_language}. Only return the translation, nothing else. NO ADDITIONAL TEXT.AND AVIOID Text to be translated,ADDITIONAL REMARKS,TRY TO understand the contexte for the translation, and avoid literal translation. " + ) + # print(text) + # print(result.data.translated_text) + return result.data.translated_text + + except Exception as e: + logging.error(f"Translation error with Ollama: {e}") + return text # Return original text on error + async def translate_text(translator, text, target_language): """Translate text to target language""" try: @@ -195,7 +239,8 @@ async def process_table_xml_safely(zip_path, target_language, translator, transl if header_text in translated_cache: column.attrib['name'] = translated_cache[header_text] else: - translated_header = await translate_text(translator, header_text, target_language) + # translated_header = await translate_text(translator, header_text, target_language) + translated_header = await translate_with_ollama(header_text, target_language) column.attrib['name'] = translated_header translated_cache[header_text] = translated_header @@ -264,7 +309,8 @@ async def translate_excel(file_path: str, target_language: str): if original_text in translated_cache: cell.value = translated_cache[original_text] else: - translated_text = await translate_text(translator, original_text, target_language) + # translated_text = await translate_text(translator, original_text, target_language) + translated_text = await translate_with_ollama(original_text, target_language) cell.value = translated_text translated_cache[original_text] = translated_text