MCP Server Replicate

Instructions

Implementation Reference

• src/mcp_server_replicate/tools/prediction_tools.py:44-60 (handler)

  The main asynchronous handler function that executes the get_prediction tool. It retrieves the prediction status from the ReplicateClient and constructs a Prediction object.

  async def get_prediction(prediction_id: str) -> Prediction:
      """Get the current status and results of a prediction.
      
      Args:
          prediction_id: The ID of the prediction to retrieve
          
      Returns:
          Prediction object containing the current status and results
          
      Raises:
          RuntimeError: If the Replicate client fails to initialize
          ValueError: If the prediction is not found
          Exception: If the status check fails
      """
      async with ReplicateClient() as client:
          result = client.get_prediction_status(prediction_id)
          return Prediction(**result)

• src/mcp_server_replicate/tools/prediction_tools.py:40-43 (registration)

  The @mcp.tool decorator that registers the get_prediction function as an MCP tool with the specified name and description.

  @mcp.tool(
      name="get_prediction",
      description="Get the current status and results of a prediction.",
  )

• src/mcp_server_replicate/models/prediction.py:28-42 (schema)

  Pydantic BaseModel defining the output schema for the get_prediction tool response.

  class Prediction(BaseModel):
      """A prediction (model run) on Replicate."""
      id: str = Field(..., description="Unique identifier for this prediction")
      version: str = Field(..., description="Model version used for this prediction")
      status: PredictionStatus = Field(..., description="Current status of the prediction")
      input: Dict[str, Any] = Field(..., description="Input parameters used for the prediction")
      output: Optional[Any] = Field(None, description="Output from the prediction if completed")
      error: Optional[str] = Field(None, description="Error message if prediction failed")
      logs: Optional[str] = Field(None, description="Execution logs from the prediction")
      created_at: datetime
      started_at: Optional[datetime] = None
      completed_at: Optional[datetime] = None
      urls: Dict[str, str] = Field(..., description="Related API URLs for this prediction")
      metrics: Optional[Dict[str, float]] = Field(None, description="Performance metrics if available")
      stream_url: Optional[str] = Field(None, description="URL for streaming output if requested") 

• src/mcp_server_replicate/replicate_client.py:401-441 (helper)

  Helper method in ReplicateClient that fetches and formats the raw prediction data from the Replicate API, used by the get_prediction tool.

  def get_prediction_status(self, prediction_id: str) -> dict[str, Any]:
      """Get the status of a prediction.
  
      Args:
          prediction_id: ID of the prediction to check
  
      Returns:
          Dict containing current status and output of the prediction
  
      Raises:
          ValueError: If the prediction is not found
          Exception: If the API request fails
      """
      if not self.client:
          raise RuntimeError("Client not initialized. Check error property for details.")
  
      try:
          # Get prediction
          prediction = self.client.predictions.get(prediction_id)
          if not prediction:
              raise ValueError(f"Prediction not found: {prediction_id}")
  
          # Return prediction status and output
          return {
              "id": prediction.id,
              "status": prediction.status,
              "output": prediction.output,
              "error": prediction.error,
              "created_at": prediction.created_at.isoformat() if prediction.created_at else None,
              "started_at": prediction.started_at.isoformat() if prediction.started_at else None,
              "completed_at": prediction.completed_at.isoformat() if prediction.completed_at else None,
              "urls": prediction.urls,
              "metrics": prediction.metrics,
          }
  
      except ValueError as err:
          logger.error(f"Validation error: {str(err)}")
          raise
      except Exception as err:
          logger.error(f"Failed to get prediction status: {str(err)}")
          raise Exception(f"Failed to get prediction status: {str(err)}") from err