Method poll_prediction

core/src/llm/replicate.rs:352–420

Poll for prediction completion

(
        &self,
        prediction_id: &str,
    )

impl ReplicateProvider {
    /// Poll for prediction completion
    async fn poll_prediction(
        &self,
        prediction_id: &str,
    ) -> GraphBitResult<ReplicatePredictionResponse> {
        let url = format!("{}/predictions/{}", self.base_url, prediction_id);
        let mut attempts = 0;
        const MAX_ATTEMPTS: u32 = 60; // 5 minutes with 5-second intervals

        loop {
            tokio::time::sleep(Duration::from_secs(5)).await;
            attempts += 1;

            let response = self
                .client
                .get(&url)
                .header("Authorization", format!("Token {}", self.api_key))
                .send()
                .await
                .map_err(|e| {
                    GraphBitError::llm_provider("replicate", format!("Polling request failed: {e}"))
                })?;

            if !response.status().is_success() {
                let error_text = response
                    .text()
                    .await
                    .unwrap_or_else(|_| "Unknown error".to_string());
                return Err(GraphBitError::llm_provider(
                    "replicate",
                    format!("Polling API error: {error_text}"),
                ));
            }

            let prediction: ReplicatePredictionResponse = response.json().await.map_err(|e| {
                GraphBitError::llm_provider(
                    "replicate",
                    format!("Failed to parse polling response: {e}"),
                )
            })?;

            match prediction.status.as_str() {
                "succeeded" => return Ok(prediction),
                "failed" | "canceled" => {
                    let error_msg = prediction.error.unwrap_or_else(|| {
                        format!(
                            "Prediction {} with status: {}",
                            prediction.status, prediction.status
                        )
                    });
                    return Err(GraphBitError::llm_provider("replicate", error_msg));
                }
                "starting" | "processing" => {
                    if attempts >= MAX_ATTEMPTS {
                        return Err(GraphBitError::llm_provider(
                            "replicate",
                            "Prediction timed out after 5 minutes".to_string(),
                        ));
                    }