""" OpenAI client to test rdrag/server.py """
import sys
from openai import OpenAI
PROMPT = "Which is bigger, a virus or a bacterium?"
BASE_URL = (["https://amd16.tail9f615d.ts.net:10000"] + sys.argv[1:])[-1]
print(BASE_URL)


def saar_berkovich(prompt=PROMPT):
    client = OpenAI(
        api_key="fake-api-key",
        base_url=BASE_URL,
        # base_url="http://localhost:8000" # change the default port if needed
    )

    # call API
    response = client.chat.completions.create(
        model="privatera",
        messages=[
            {
                "role": "user",
                "content": prompt,
            }
        ],
    )

    print(response.choices[0].message.content)
    return response


def virus(prompt=PROMPT):
    client = OpenAI(base_url=LLM_BASE_URL, api_key="sk-xxx")
    response = client.chat.completions.create(
        model='lmstudio-ai_gemma-2b-it',
        messages=[
            {
                "role": "user",
                "content": prompt,
            }
        ],
    )

    print('## LLM RESPONSE OBJECT:')
    print(response)

    response_text = response.choices[0].message.content
    return response_text

    # # print the response from the chatbot in real-time
    # print('Chatbot response:')
    # for chunk in stream:
    #     print(chunk['message']['content'], end='', flush=True)
