src/llm/llm.py

from langchain_openai import ChatOpenAI


class OpenAILLM:
    def __init__(self, model_name) -> None:
        self.oai_llm = ChatOpenAI(model=model_name)

    def run(self, query: str):
        return self.oai_llm.invoke(query)