utils/gpt.py

import cohere


def test_api_key(api_key: str):

    try:
        # try to just generate 3 tokens
        co = cohere.Client(
            api_key=api_key,
        )
        response = co.generate(prompt="sample prompt", max_tokens=3)
        return True
    except:
        return False


def gpt_stream_response(prompt: str, api_key: str):
    """Get response from Cohere and stream response"""

    co = cohere.Client(
        api_key=api_key,
    )

    stream = co.chat_stream(message=prompt)

    for event in stream:
        if event.event_type == "text-generation":
            yield event.text


def gpt_response(prompt: str, api_key: str) -> str:
    """Get response from Cohere, with option to get output in json format"""

    co = cohere.Client(
        api_key=api_key,
    )

    response = co.chat(message=prompt)

    return response.text