translator/llm_client.py

from huggingface_hub import InferenceClient
from openai import OpenAI
import os

HF_MODEL_DEFAULT = "mistralai/Mistral-7B-Instruct-v0.3"
OPENAI_MODEL_DEFAULT = "gpt-4o-mini"


def generate_with_hf(prompt, model=HF_MODEL_DEFAULT):
    client = InferenceClient(model=model, token=os.getenv("HF_TOKEN"))
    response = client.chat_completion(
        messages=[
            {"role": "system", "content": "You are a Unity→Unreal translation expert."},
            {"role": "user", "content": prompt},
        ],
        max_tokens=1024,
        temperature=0.7,
    )
    return response.choices[0].message["content"]

def generate_with_openai(prompt, model=OPENAI_MODEL_DEFAULT):
    client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
    completion = client.chat.completions.create(
        model=model,
        messages=[
            {"role": "system", "content": "You are a Unity→Unreal translation expert."},
            {"role": "user", "content": prompt},
        ],
    )
    return completion.choices[0].message.content