openai_model.py

# ruff: noqa: F401

import os
import sys
from platform import node

import rich
from loguru import logger
from smolagents import LiteLLMModel, OpenAIServerModel

from exit_gracefully import exit_gracefully

print = rich.get_console().print  # noqa


def openai_model(
    model_id = None,
    api_base = None,
    api_key = None,
    **kwargs,
):
    kwargs = kwargs or {}
    # default siliconflow
    # api_base = api_base or "https://api.siliconflow.cn/v1"
    # api_key = api_key or os.getenv("SILICONFLOW_API_KEY")
    # model_id = model_id or "deepseek-ai/DeepSeek-V3"

    # default llama4
    api_base = api_base or "https://api.llama.com/compat/v1"

    if "golay" in node() and ("llama.com" in api_base or "openai.com" in api_base):
        os.environ.update(
            HTTPS_PROXY="http://localhost:8081",
            HTTP_PROXY="http://localhost:8081",
            ALL_PROXY="http://localhost:8081",
            NO_PROXY="localhost,127.0.0.1",
        )

    api_key = api_key or os.getenv("LLAMA_API_KEY")
    if isinstance(api_key, str):
        # LLAMA_API_KEY contains | and in win10 need to assign env var with ""
        api_key = api_key.strip('"')
    assert api_key, "LLAMA_API_KEY not set, set it and try again"

    default =  "Llama-4-Maverick-17B-128E-Instruct-FP8"
    # default = "Llama-4-Scout-17B-16E-Instruct-FP8"
    logger.debug(f"{default=}")

    model_id = model_id or default

    return OpenAIServerModel(
        model_id,
        api_base=api_base,
        api_key=api_key,
        # temperature=0.,
        **kwargs,
    )

def main():
    messages = [{'role': 'user', 'content': 'Say this is a test.'}]
    logger.debug(sys.argv)

    if not sys.argv[1:]:
        model = openai_model()
        logger.debug(model(messages))
        return

    if len(sys.argv[1:]) < 3:
        raise SystemExit("Provide at least three args (model_id, api_base, api_key)")

    model_id, api_base, api_key, *_ = sys.argv[1:]
    model = openai_model(model_id, api_base, api_key)
    try:
        response = model(messages)
        logger.debug(response)
    except Exception as e:
        logger.error(e)
        return
    try:
        print(response.content)
    except Exception as e:
        logger.error(e)


if __name__ == "__main__":
    main()
    # python openai_model.py
    #  deepseek-ai/DeepSeek-V3 https://api.siliconflow.cn/v1 %SILICONFLOW_API_KEY%

    # python openai_model.py  grok-3-beta https://api.x.ai/v1 %XAI_API_KEY%

    #   gemini-2.5-flash-preview-04-17  https://generativelanguage.googleapis.com/v1beta %GEMINI_API_KEY%
    #   gemini-2.0-flash

    # https://api.together.ai/models/deepseek-ai/DeepSeek-V3
    # deepseek-ai/DeepSeek-V3 https://api.together.xyz/v1 %TOGETHER_API_KEY%

    #  deepseek-chat  https://litellm.dattw.eu.org/v1 %LITELLM_API_KEY%

    #

    # LLM API proxy: https://linux.do/t/topic/290871