# Create and wrap the OpenAI client with Helicone's base URL
client = OpenAI(
    api_key=os.environ["HELICONE_API_KEY"],
    base_url="https://ai-gateway.helicone.ai"
)
client = track_openai(client)

@track
def generate_response(prompt: str):
    response = client.chat.completions.create(
        model="gpt-3.5-turbo",
        messages=[
            {"role": "system", "content": "You are a knowledgeable AI assistant."},
            {"role": "user", "content": prompt}
        ]
    )
    return response.choices[0].message.content

@track
def refine_response(initial_response: str):
    response = client.chat.completions.create(
        model="gpt-4",
        messages=[
            {"role": "system", "content": "You enhance and polish text responses."},
            {"role": "user", "content": f"Please improve this response: {initial_response}"}
        ]
    )
    return response.choices[0].message.content

@track(project_name="helicone-integration-demo")
def generate_and_refine(prompt: str):
    # First LLM call: Generate initial response
    initial = generate_response(prompt)
    
    # Second LLM call: Refine the response
    refined = refine_response(initial)
    
    return refined