from openworm_ai.utils.llms import LLM_GPT4o
from openworm_ai.utils.llms import LLM_LLAMA2
from openworm_ai.utils.llms import LLM_GEMINI_2F
from openworm_ai.utils.llms import LLM_GEMINI_25F
from openworm_ai.utils.llms import LLM_COHERE


from ring import generate_response

import sys

question = "What is the most common type of neuron in the brain?"

llm_ver = LLM_GPT4o


if "-g2" in sys.argv:
    llm_ver = LLM_GEMINI_2F

if "-g25" in sys.argv:
    llm_ver = LLM_GEMINI_25F

if "-l" in sys.argv:
    llm_ver = LLM_LLAMA2

if "-co" in sys.argv:
    llm_ver = LLM_COHERE


print("--------------------------------------------------------")
print("Asking question:\n   %s" % question)
print("--------------------------------------------------------")

print(" ... Connecting to: %s" % llm_ver)

response = generate_response(question, llm_ver, temperature=0, only_celegans=False)


print("--------------------------------------------------------")
print("Answer:\n   %s" % response)
print("--------------------------------------------------------")
print()