trace_decorator.py

trace_decorator.py

@trace()
def _get_openai_client():
    return openai.Client()


@trace()
def _query_openai(
    prompt: str,
    model: str = "gpt-4",
    chat_history: Optional[list] = None
):
    chat_history = chat_history or []
    client = _get_openai_client()
    result = client.chat.completions.create(
        model=model,
        messages=[
            {"role": "system", "content": "You are a helpful assistant"},
            *chat_history,
            {"role": "user", "content": prompt},
        ],
    )
    return result.choices[0].message.content