llms.py

import re
import os
import json
import logging
from dotenv import load_dotenv

from openai import OpenAI
from tenacity import retry, stop_after_attempt, wait_random_exponential, before_sleep_log


load_dotenv()

logger = logging.getLogger(__name__)

# Define retry strategy parameters
RETRY_TIMES = int(os.getenv('RETRY_TIMES'))
WAIT_TIME_LOWER = int(os.getenv('WAIT_TIME_LOWER'))
WAIT_TIME_UPPER = int(os.getenv('WAIT_TIME_UPPER'))

OPENAI_BASE_URL = os.getenv('OPENAI_BASE_URL')
OPENAI_API_KEY = os.getenv('OPENAI_API_KEY')
MODEL = os.getenv('OPENAI_MODEL')


common_params = {}

if os.getenv('OPENAI_MAX_TOKENS'):
    common_params["max_tokens"] = int(os.getenv('OPENAI_MAX_TOKENS'))

if os.getenv('OPENAI_TEMPERATURE'):
    common_params["temperature"] = float(os.getenv('OPENAI_TEMPERATURE'))

if os.getenv('OPENAI_TIMEOUT'):
    common_params["timeout"] = int(os.getenv('OPENAI_TIMEOUT'))


client = OpenAI(
    base_url=OPENAI_BASE_URL,
    api_key=OPENAI_API_KEY
)


@retry(
    wait=wait_random_exponential(min=WAIT_TIME_LOWER, max=WAIT_TIME_UPPER), 
    stop=stop_after_attempt(RETRY_TIMES), 
    reraise=True,
    before_sleep=before_sleep_log(logger, logging.WARNING)
)
def llm_request(prompt):
    """
    Sends a request to the specified language model with a given prompt.

    Args:
        prompt (str): The input text or message to send to the model.

    Returns:
        str: The response generated by the model.
    """

    response_obj = client.chat.completions.create(
        model=MODEL,
        messages=[
            {
                'role': 'user', 
                'content': prompt
            }
        ],
        **common_params
    )

    return response_obj.choices[0].message.content

@retry(
    wait=wait_random_exponential(min=WAIT_TIME_LOWER, max=WAIT_TIME_UPPER),
    stop=stop_after_attempt(RETRY_TIMES),
    reraise=True,
    before_sleep=before_sleep_log(logger, logging.WARNING)
)
def llm_request_for_json(prompt):

    response_obj = client.chat.completions.create(
        model=MODEL,
        messages=[{'role': 'user', 'content': prompt}],
        **common_params
    )

    content = response_obj.choices[0].message.content or ""

    match = re.search(r"```json\s*(\{.*?\})\s*```", content, re.DOTALL)
    if not match:
        raise ValueError(f"No JSON block found in model output: {content}")

    json_str = match.group(1).strip()

    return json.loads(json_str)


if __name__ == '__main__':
    r = llm_request_for_json('hello')
    print(r)