modules/translator.py

import os
import gradio as gr
from langchain_core.pydantic_v1 import BaseModel, Field
from langchain.prompts import HumanMessagePromptTemplate, ChatPromptTemplate
from langchain.output_parsers import PydanticOutputParser
from langchain_openai import ChatOpenAI

chat = ChatOpenAI()

# Define the Pydantic Model
class TextTranslator(BaseModel):
    output: str = Field(description="Python string containing the output text translated in the desired language")
    
output_parser = PydanticOutputParser(pydantic_object=TextTranslator)
format_instructions = output_parser.get_format_instructions()

def text_translator(input_text: str, language: str) -> str:
    human_template = """Enter the text that you want to translate: 
                {input_text}, and enter the language that you want it to translate to {language}. {format_instructions}"""
    human_message_prompt = HumanMessagePromptTemplate.from_template(human_template)
    
    chat_prompt = ChatPromptTemplate.from_messages([human_message_prompt])
    
    prompt = chat_prompt.format_prompt(input_text=input_text, language=language, format_instructions=format_instructions)
    
    messages = prompt.to_messages()
    
    response = chat(messages=messages)
    
    output = output_parser.parse(response.content)
    
    output_text = output.output
    
    return output_text

# แยก UI เป็นฟังก์ชัน เพื่อนำไปใส่ใน Gradio Tab ได้
def text_translator_ui():
    gr.Markdown("### Text Translator\nTranslate text into any language using AI.")
    
    input_text = gr.Textbox(label="Enter the text that you want to translate")
    input_lang = gr.Textbox(label="Enter the language that you want it to translate to", placeholder="Example: Hindi, French, Bengali, etc.")
    output_text = gr.Textbox(label="Translated text")
    
    translate_button = gr.Button("Translate")
    translate_button.click(fn=text_translator, inputs=[input_text, input_lang], outputs=output_text)