# Configure Opik integration for LLM calls within steps
set_global_handler("opik")

class QueryEvent(Event):
    """Event for passing query through workflow."""
    query: str

class MyRAGWorkflow(Workflow):
    """Simple RAG workflow with tracked steps."""
    
    @step
    @opik.track()
    async def retrieve_context(self, ev: StartEvent) -> QueryEvent:
        """Retrieve relevant context for the query."""
        query = ev.get("query", "")
        # Your retrieval logic here
        context = f"Context for: {query}"
        return QueryEvent(query=f"{context} | {query}")
    
    @step
    @opik.track()
    async def generate_response(self, ev: QueryEvent) -> StopEvent:
        """Generate final response using the context."""
        # Your generation logic here
        result = f"Response based on: {ev.query}"
        return StopEvent(result=result)