# 3. Create our LLM agent within the parent span context
with llm_config:
    my_agent = ConversableAgent(
        name="helpful_agent",
        system_message="You are a poetic AI assistant, respond in rhyme.",
    )


def main(message):
    response = my_agent.run(message=message, max_turns=2, user_input=True)

    # 5. Iterate through the chat automatically with console output
    response.process()

    # 6. Print the chat
    print(response.messages)

    return response.messages


if __name__ == "__main__":
    tracer, provider = setup_telemetry()

    # 4. Run the agent with a prompt
    with tracer.start_as_current_span(my_agent.name) as agent_span:
        message = "In one sentence, what's the big deal about AI?"

        agent_span.set_attribute("input", message)  # Manually log the question

        response = main(message)

        # Manually log the response
        agent_span.set_attribute("output", response)

    # Force flush all spans to ensure they are exported
    provider = trace.get_tracer_provider()
    provider.force_flush()
```

## Further improvements

If you would like to see us improve this integration, simply open a new feature
request on [Github](https://github.com/comet-ml/opik/issues).