llm.py

from openai import OpenAI
import json
from dotenv import load_dotenv
import os

load_dotenv()

LLM_URL = os.getenv("LLM_URL", "https://generativelanguage.googleapis.com/v1beta/openai/")
LLM_API_KEY = os.getenv("LLM_API_KEY")
LLM_MODEL = os.getenv("LLM_MODEL")
LLM_MODEL_LEARN_LM = os.getenv("LLM_MODEL_LEARN_LM")

client = OpenAI(
    api_key=LLM_API_KEY,
    base_url=LLM_URL
)

def llm_request(system_prompt, user_prompt):
    response = client.chat.completions.create(
        model=LLM_MODEL_LEARN_LM,
        n=1,
        messages=[
            {"role": "user", "content": system_prompt},
            {"role": "user", "content": user_prompt}
        ],
        temperature=0.8
    )
    return response.choices[0].message.content


if __name__ == "__main__":
    print(f"url: {LLM_URL}")
    print(f"key: {LLM_API_KEY}")
    print(f"model: {LLM_MODEL}")
    print(f"model: {LLM_MODEL_LEARN_LM}")

    res = llm_request(
        system_prompt="Bạn là một chuyên gia toán học, giải toán và đưa ra lời giải ở định dạng Latex",
        user_prompt="Tìm $x$ biết $2x+1=0$"
    )

    print(res)