Parse prediction response into LlmResponse
(
&self,
prediction: ReplicatePredictionResponse,
)
| 421 | |
| 422 | /// Parse prediction response into LlmResponse |
| 423 | fn parse_prediction_response( |
| 424 | &self, |
| 425 | prediction: ReplicatePredictionResponse, |
| 426 | ) -> GraphBitResult<LlmResponse> { |
| 427 | if prediction.status != "succeeded" { |
| 428 | let error_msg = prediction |
| 429 | .error |
| 430 | .unwrap_or_else(|| format!("Prediction failed with status: {}", prediction.status)); |
| 431 | return Err(GraphBitError::llm_provider("replicate", error_msg)); |
| 432 | } |
| 433 | |
| 434 | let output = prediction.output.ok_or_else(|| { |
| 435 | GraphBitError::llm_provider( |
| 436 | "replicate", |
| 437 | "No output in successful prediction".to_string(), |
| 438 | ) |
| 439 | })?; |
| 440 | |
| 441 | // Extract content from output |
| 442 | let content = match output { |
| 443 | serde_json::Value::String(s) => s, |
| 444 | serde_json::Value::Array(arr) => { |
| 445 | // Join array elements (common for streaming models) |
| 446 | arr.iter() |
| 447 | .filter_map(|v| v.as_str()) |
| 448 | .collect::<Vec<_>>() |
| 449 | .join("") |
| 450 | } |
| 451 | _ => output.to_string(), |
| 452 | }; |
| 453 | |
| 454 | // Parse tool calls if present |
| 455 | let (clean_content, tool_calls) = if self.model_supports_function_calling() { |
| 456 | Self::parse_tool_calls_from_response(&content) |
| 457 | } else { |
| 458 | (content, Vec::new()) |
| 459 | }; |
| 460 | |
| 461 | // Create usage information |
| 462 | let usage = LlmUsage::new( |
| 463 | 0, // Replicate doesn't provide input token count |
| 464 | clean_content.split_whitespace().count() as u32, // Rough estimate |
| 465 | ); |
| 466 | |
| 467 | let finish_reason = if !tool_calls.is_empty() { |
| 468 | FinishReason::ToolCalls |
| 469 | } else { |
| 470 | FinishReason::Stop |
| 471 | }; |
| 472 | |
| 473 | let mut response = LlmResponse::new(clean_content, &self.model) |
| 474 | .with_tool_calls(tool_calls) |
| 475 | .with_usage(usage) |
| 476 | .with_finish_reason(finish_reason) |
| 477 | .with_id(prediction.id); |
| 478 | |
| 479 | // Add prediction time to metadata if available |
| 480 | if let Some(metrics) = prediction.metrics { |
no test coverage detected